# -*- coding: utf-8 -*-
"""
Created on Tue Jul 21 23:01:40 2015

@author: xhxiao
"""

import tomopy 
import os
import numpy as np
from tomopy.recon.rotation import write_center
from tomopy.recon.algorithm import recon
from scipy import misc
import glob, fnmatch
import h5py
import shutil, os

####----------------- Section 1: Input file -- Staragt -----------------  
user_top_dir = '/run/media/user2bmb/APS3/APS_2BM_2017_11'
center_path = '/run/media/user2bmb/APS4/APS_2BM_2017_11/data_center/'

#sam_idx = 1
Exp_idx = 144
#data_top_dir = user_top_dir + 's' + '{:03}'.format(sam_idx) + '*/'
data_top_dir = user_top_dir + '/'

data_dir = glob.glob(os.path.join(data_top_dir,'Exp'+'{:03}'.format(Exp_idx)+'*'))[0]
output_dir = data_dir

data_files = np.sort(fnmatch.filter(os.listdir(data_dir),'*.hdf'))
if len(data_files) == 0:
    data_files = np.sort(fnmatch.filter(os.listdir(data_dir),'*.h'))	

file_data = data_files[0]

print file_data

BlankAt = 90
if BlankAt == 0:
   first_missing_end = 87
   second_missing_start = 1396
else:
    missing_start = 657
    missing_end = 828
    open_start = 739
    open_end = 749
    missing_width = missing_end - missing_start
avg_width = 20

file_name = os.path.join(data_dir, file_data)
output_file = output_dir+'/recon_'+file_data.split(".")[-2]+'/recon_'+file_data.split(".")[-2] +'_'
####----------------- Section 1: Input file -- End -----------------  

####----------------- Section 2: Parameter configuration -- Start -----------------  
f = h5py.File(file_name,"r")
try:
    arr = f["exchange/data"]
except:
    try:
        arr = f["entry/data/data"]
    except:
        arr = f["MyDataDir/TomoShots"] 
numProj = (arr.shape)[0]
numSlices = (arr.shape)[1]
imgWidth = (arr.shape)[2]
f.close()

offset = 0    # offset from the top slice; for full volume, set to 0
#numRecSlices = 600    # For full volume, comment this out# center = 1290.0  # proj_49
numRecSlices = numSlices    # For full volume, uncomment this line
chunk_size = 200
if chunk_size > numRecSlices:
    chunk_size = numRecSlices
margin_slices = 30
num_chunk = np.int(numRecSlices/(chunk_size-margin_slices)) + 1
if numRecSlices == chunk_size:
    num_chunk = 1
    
z = 20
eng = 60
pxl = 0.65e-4
rat = 1.0e-02

zinger_level = 1000
###----------------- Section 2: Parameter configuration -- End -----------------   

######----------------- Section 3: Finding center -- Start -----------------
sliceStart = 660
sliceEnd = sliceStart+40
f = h5py.File(file_name,"r")
try:
    arr = f["exchange/data"]
except:
    try:
        arr = f["entry/data/data"]
    except:
        arr = f["MyDataDir/TomoShots"]     
data = arr[0:numProj,sliceStart:sliceEnd,:]
f.close()

f = h5py.File(file_name,"r")
try:
    arr = f["exchange/data_white"]
except:
    try:
        arr = f["entry/data/data"]
    except:
        arr = f["MyDataDir/TomoShots"] 
flat = arr[1:9,sliceStart:sliceEnd,:]
f.close()

f = h5py.File(file_name,"r")
try:
    arr = f["exchange/data_dark"]
except:
    try:
        arr = f["entry/data/data"]
    except:
        arr = f["MyDataDir/TomoShots"]     
dark = arr[1:9,sliceStart:sliceEnd,:]
f.close()

if BlankAt == 0:
    missing_width = int(first_missing_end + (numProj-second_missing_start))
    missing_start = int(int(numProj/2) - (numProj - second_missing_start) + 1)
    missing_end = int(int(numProj/2) + first_missing_end)
theta = np.linspace(0,np.pi,num=numProj)
 

# zinger_removal
data = tomopy.misc.corr.remove_outlier(data,zinger_level,size=15,axis=0)
flat = tomopy.misc.corr.remove_outlier(flat,zinger_level,size=15,axis=0)
                           
# there is modification in below normalize routine# center = 1290.5  # proj_116
data = tomopy.prep.normalize.normalize(data,flat,dark)
data = tomopy.prep.normalize.normalize_bg(data,air=10)

data_size = data.shape

if BlankAt == 0: 
   data = np.roll(data,int(numProj/2),axis=0)
   data[:int(numProj/2),:] = data[:int(numProj/2),:,::-1]
#tomopy.write_tiff(data,'/media/2BM_Backup48/debug/shift_img.tiff')

data = tomopy.prep.phase.retrieve_phase(data,pixel_size=pxl,dist=z, energy=eng,alpha=rat,pad=True)
# data = tomopy.prep.stripe.remove_stripe_fw(data,level=6,wname='sym16',sigma=2,pad=True)
data = tomopy.prep.stripe.remove_stripe_ti(data,alpha=5)
data = tomopy.prep.stripe.remove_stripe_sf(data,size=31)                                        
data = tomopy.prep.normalize.normalize_bg(data,air=10)


if os.path.isdir(center_path):
    shutil.rmtree(center_path)

center_shift = 5
center_shift_w = 15
data1 = np.empty_like(data)
print 'index,','roll_dist,','center'  
for jj in range(int(center_shift_w/0.5)):
    roll_dist = int(2*(center_shift+jj*0.5)+1)
    print jj,', ',roll_dist,', ',data_size[2]/2+(center_shift+jj*0.5) 
    data1[:] = data 
    if BlankAt == 0:
        data1[:int(numProj/2),:] =np.roll(data1[:int(numProj/2),:],roll_dist,axis=2)
    for ii in range(missing_start,open_start):
        data1[ii,:,:] = np.roll(np.mean(np.roll(data1[(ii - missing_width - avg_width):(ii - missing_width),:,:],roll_dist,axis=2),axis=0) + \
                            np.mean(np.roll(data1[(ii + missing_width):(ii + missing_width + avg_width),:,:],roll_dist,axis=2),axis=0),-1*roll_dist)/2.0
    for ii in range(open_end,missing_end):
        data1[ii,:,:] = np.roll(np.mean(np.roll(data1[(ii - missing_width - avg_width):(ii - missing_width),:,:],roll_dist,axis=2),axis=0) + \
                            np.mean(np.roll(data1[(ii + missing_width):(ii + missing_width + avg_width),:,:],roll_dist,axis=2),axis=0),-1*roll_dist)/2.0
    
    data1[open_start:open_end,:,:] = data[open_start:open_end,:,:]               
    data1 = tomopy.prep.normalize.minus_log(data1)
    write_center(data1[:,9:11,:], theta, dpath=center_path, 
                 cen_range=(data_size[2]/2+(center_shift+jj*0.5),data_size[2]/2+(center_shift+(jj+1)*0.5),0.5),overwrite=False)
######----------------- Section 3: Finding center -- End -----------------

# center = 1289.0  # proj_3
# center = 1291.5  # proj_4
# center = 1290.5  # proj_5
# center = 1293.5  # proj_6
# center = 1293.0  # proj_7
# center = 1291.5  # proj_8
# center = 1297.5  # proj_9
# center = 1297.5  # proj_10
# center = 1290.5  # proj_11
# center = 1285.5  # proj_12
# center = 1285.0  # proj_13
# center = 1286.0  # proj_14
# center = 1287.0  # proj_15
# center = 1284.0  # proj_16
# center = 1295.0  # proj_17
# center = 1286.5  # proj_18
# center = 1287.5  # proj_19
# center = 1286.5  # proj_20
# center = 1286.0  # proj_21
# center = 1286.5  # proj_22
# center = 1288.5  # proj_23
# center = 1288.5  # proj_24
# center = 1285.0  # proj_25
# center = 1289.0  # proj_26
# center = 1290.5  # proj_27
# center = 1289.5  # proj_28
# center = 1286.0  # proj_29
# center = 1283.0  # proj_30
# center = 1286.5  # proj_31
# center = 1287.0  # proj_32
# center = 1296.5  # proj_33
# center = 1296.0  # proj_34
# center = 1286.0  # proj_35
# center = 1286.0  # proj_36
# center = 1285.0  # proj_37
# center = 1290.0  # proj_38
# center = 1293.5  # proj_39
# center = 1290.0  # proj_40
# center = 1290.0  # proj_41
# center = 1284.5  # proj_42
# center =  # proj_43  # empty
# center = 1290.0  # proj_44
# center = 1286.0  # proj_45
# center = 1286.0  # proj_46
# center = 1286.0  # proj_47
# center = 1287.5  # proj_48
# center = 1290.0  # proj_491287
# center = 1290.5  # proj_50
# center = 1290.0  # proj_51
# center = 1297.5  # proj_52
# center = 1289.0  # proj_53
# center = 1285.0  # proj_54
# center = 1285.0  # proj_55
# center = 1286.0  # proj_56
# center = 1288.5  # proj_57
# center = 1290.0  # proj_58
# center = 1296.5  # proj_59
# center = 1295.0  # proj_60
# center = 1294.5  # proj_61
# center = 1300.5  # proj_62
# center = 1287.0  # proj_63
# center = 1290.5  # proj_64
# center = 1299.0  # proj_65
# center = 1287.0  # proj_66
# center = 1288.5  # proj_67
# center = 1299.0  # proj_68
# center = 1286.5  # proj_69
# center = 1286.5  # proj_70
# center = 1286.5  # proj_71
# center = 1294.5  # proj_72
# center = 1287.5  # proj_73
# center = 1298.0  # proj_74
# center = 1285.0  # proj_75
# center = 1290.0  # proj_76
# center = 1299.5  # proj_77
# center = 1289.5  # proj_78
# center = 1289.5  # proj_79
# center = 1300.5  # proj_80
# center = 1291.5  # proj_81
# center = 1292.0  # proj_82
# center = 1292.0  # proj_83
# center = 1289.0  # proj_84
# center = 1290.5  # proj_85
# center =         # proj_86 "Mean of zero slice" warning, black slices
# center = 1288.5  # proj_87
# center = 1291.0  # proj_88
# center = 1298.0  # proj_89
# center = 1288.5  # proj_90
# center = 1291.0  # proj_91
# center = 1289.5  # proj_92
# center = 1294.5  # proj_93
# center = 1291.5  # proj_94
# center = 1299.0  # proj_95
# center = 1290.5  # proj_96
# center = 1287.5  # proj_97
# center = 1296.5  # proj_98
# center = 1289.0  # proj_99
# center = 1289.0  # proj_100
# center = 1286.0  # proj_101
# center = 1293.0  # proj_102
# center = 1288.0  # proj_103
# center = 1295.5  # proj_104
# center = 1295.0  # proj_105
# center = 1290.5  # proj_106
# center = 1287.5  # proj_107
# center = 1297.5  # proj_108
# center = 1288.0  # proj_109
# center = 1284.0  # proj_110
# center = 1298.5  # proj_111
# center = 1290.0  # proj_112
# center = 1300.0  # proj_113
# center = 1288.5  # proj_114
# center = 1284.5  # proj_115
# center = 1290.5  # proj_116
# center = 1293.5  # proj_117
# center = 1281.5  # proj_118
# center = 1287.5  # proj_119
# center = 1293.5  # proj_120
# center = 1292.0  # proj_121
# center = 1289.0  # proj_122
# center = 1290.0  # proj_123
# center = 1282.5  # proj_124
# center = 1287.0  # proj_125
# center = 1295.5  # proj_126
# center = 1285.5  # proj_127
# center = 1291.0  # proj_128
# center = 1284.0  # proj_129
# center = 1296.5  # proj_130
# center = 1293.5  # proj_131
# center = 1299.0  # proj_132
# center = 1289.5  # proj_133
# center = 1287.5  # proj_134
# center = 1288.0  # proj_135
# center = 1288.0  # proj_136
# center = 1291.0  # proj_137
# center = 1286.0  # proj_138
# center = 1292.0  # proj_139
# center = 1287.0  # proj_140
# center = 1298.5  # proj_141
# center = 1290.5  # proj_142
# center = 1288.0  # proj_143
# center = 1286.0  # proj_144


########----------------- Section 4: Full reconstruction -- Start -----------------
#for ii in xrange(num_chunk):
#    if ii == 0:
#        sliceStart = offset + ii*chunk_size
#        sliceEnd = offset + (ii+1)*chunk_size
#    else:
#        sliceStart = offset + ii*(chunk_size-margin_slices)
#        sliceEnd = offset + sliceStart + chunk_size
#        if sliceEnd > (offset+numRecSlices):
#            sliceEnd = offset+numRecSlices
#        if sliceEnd > numSlices:
#            sliceEnd = numSlices
#    
#    if (sliceEnd - sliceStart) <= margin_slices:
#        print 'Reconstruction finishes!'
#        break        
#
#    f = h5py.File(file_name,"r")
#    try:
#        arr = f["exchange/data"]
#    except:
#        try:
#            arr = f["MyDataDir/TomoShots"]
#        except:
#            arr = f["entry/data/data"]
#    data = arr[0:numProj,sliceStart:sliceEnd,:]
#    f.close()
#    
#    f = h5py.File(file_name,"r")
#    try:
#        arr = f["exchange/data_white"]
#    except:
#        try:
#            arr = f["MyDataDir/TomoShots"]
#        except:
#            arr = f["entry/data/data"]
#    white = arr[1:9,sliceStart:sliceEnd,:]
#    f.close()
#    
#    f = h5py.File(file_name,"r")
#    try:
#        arr = f["exchange/data_dark"]
#    except:
#        try:
#            arr = f["MyDataDir/TomoShots"]
#        except:
#            arr = f["entry/data/data"]
#    dark = arr[1:9,sliceStart:sliceEnd,:]
#    f.close()
#    data_size = data.shape
#    theta = np.linspace(0,np.pi,num=data_size[0]) 
#    print 'data is read'
#    print ii   
#    roll_dist = int(2.0*(center - data_size[2]/2.0)+1)
#    
#    # normalize projection images; for now you need to do below two operations in sequence
#    data = tomopy.prep.normalize.normalize(data,white,dark)
#    data = tomopy.prep.normalize.normalize_bg(data,air=10)
#
#    # phase retrieval
#    data = tomopy.prep.phase.retrieve_phase(data,pixel_size=pxl,dist=z,energy=eng,alpha=rat,pad=True) 
#
#    # remove stripes in sinograms
#    data = tomopy.prep.stripe.remove_stripe_ti(data,alpha=5)
#    data = tomopy.prep.stripe.remove_stripe_sf(data,size=31)                                        
#    data = tomopy.prep.normalize.normalize_bg(data,air=10)
#    
#    for ii in range(missing_start,open_start):
#        data[ii,:,:] = np.roll(np.mean(np.roll(data[(ii - missing_width - avg_width):(ii - missing_width),:,:],roll_dist,axis=2),axis=0) + \
#                               np.mean(np.roll(data[(ii + missing_width):(ii + missing_width + avg_width),:,:],roll_dist,axis=2),axis=0),-1*roll_dist)/2.0
#    for ii in range(open_end,missing_end):
#        data[ii,:,:] = np.roll(np.mean(np.roll(data[(ii - missing_width - avg_width):(ii - missing_width),:,:],roll_dist,axis=2),axis=0) + \
#                               np.mean(np.roll(data[(ii + missing_width):(ii + missing_width + avg_width),:,:],roll_dist,axis=2),axis=0),-1*roll_dist)/2.0
#                                             
#    # tomo reconstruction
#    data = tomopy.prep.normalize.minus_log(data)                                         
#    data_recon = recon(data,theta,center=center,algorithm='gridrec')
#    
#    # save reconstructions
#    tomopy.io.writer.write_tiff_stack(data_recon[np.int(margin_slices/2):(sliceEnd-sliceStart-np.int(margin_slices/2)),:,:], 
#                                      axis = 0,
#                                      fname = output_file, 
#                                      start = sliceStart+np.int(margin_slices/2),
#                                      overwrite = True)
##########----------------- Section 4: Full reconstruction -- End -----------------
#
#