import math
import numpy as np
import os
import h5py
from lib.util import *
from lib.generateConstants import *
from lib.reduce import *
import sys
import subprocess
import pathlib
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import xarray as xr
from scipy.optimize import curve_fit
plt.rcParams.update({'font.size':10})

class Config:

    def __init__(self, size_xy: int = None, size_z: int = None, phys_size: float = None,
                a_nm: float = None, b_nm: float = None, theta: float = None, radius_nm: float = None):

        # Simulation Parameters
        self.size_xy = size_xy
        self.size_z  = size_z
        self.phys_size   = phys_size

        # Materials Parameters
        self.num_materials = 2

        # Lattice Parameters
        self.a_nm  = a_nm   # Unit cell length a in nm
        self.b_nm  = b_nm   # Unit cell length b in nm
        self.theta = theta  # Angle between the unit cell edges a and b
        self.radius_nm = radius_nm    # Radius of the cylinders in nm

class Morphology:
    def __init__(self, config: Config = None):
        self.cfg = config
        self.box = None
        self.coords = None
        self.circle_tile = None

    def makeMorphology(self):
        # Initialize box
        self.box = generateBox(self.cfg)
        # Generate lattice coordinates
        self.coords = generateCoords(self.cfg)
        # Populate the box
        self.box = populateBox(self)
        # Generate the material maps that will be used for the Cy-RSoXS input
        self.unaligned, self.aligned = generateMaterialMaps(self.box, self.cfg.num_materials)

def generateBox(cfg):
    '''Generate the simulation box.'''
    box = np.zeros((cfg.size_xy, cfg.size_xy, cfg.size_z, cfg.num_materials + 1))
    box[:, :, :, 0] = 1
    return box

def generateCoords(cfg):

    def calcRange(size_xy, a_vector, b_vector, angle, range_overshoot):
        cosang = np.cos(np.deg2rad(angle))
        sinang = np.sin(np.deg2rad(angle))

        m_range = int(round(range_overshoot*(size_xy*cosang + size_xy*sinang)/np.amax(np.absolute(b_vector))))
        n_range = int(round(range_overshoot*(size_xy*sinang + size_xy*cosang)/np.amax(np.absolute(a_vector))))

        return m_range, n_range

    # Convert a and b lattice edge lengths into voxel lengths (integers) based on the phys_size
    a_voxels = int(round(cfg.a_nm/cfg.phys_size))
    b_voxels = int(round(cfg.b_nm/cfg.phys_size))

    # The lattice point coordinates are generated by looping over integer multiples of the lattice vectors (n*a + m*b)
    # The ranges of n and m are estimated by dividing the simulation box size by the a and b length scales
    # Lattices with non-90deg angles between a and b require overshooting the typical counting to fill the simulation box
    coords = []
    range_overshoot = 1.1

    b_0 = b_voxels*np.cos(np.deg2rad(cfg.theta))
    b_1 = b_voxels*np.sin(np.deg2rad(cfg.theta))

    a_vector = [a_voxels, 0]
    b_vector = [b_0, b_1]

    angle = 0
    origin = int(round(cfg.size_xy/2))

    m_range, n_range = calcRange(cfg.size_xy, a_vector, b_vector, angle, range_overshoot)

    # Coordinates are generated from -n to n and -m to m
    # Only coordinates inside the simulation box are added to the list
    for m in range(-m_range, m_range + 1):
        for n in range(-n_range, n_range + 1):

            new_coord_x = int(round(origin + n*a_vector[0] + m*b_vector[0]))
            new_coord_y = int(round(origin + n*a_vector[1] + m*b_vector[1]))

            if new_coord_x in range(0, cfg.size_xy) and new_coord_y in range(0, cfg.size_xy):
                coords.append([new_coord_x, new_coord_y, 1])

    return coords

def generateCircleTile(cfg):

    def getMajorityMaterialID(fractions):
        max_fraction = np.amax(fractions)
        material_ID  = np.where(fractions == max_fraction)
        return material_ID[0]

    radius = int(round(cfg.radius_nm/cfg.phys_size))
    width = int(round(2 * radius))
    circle_tile = np.zeros((width, width, cfg.num_materials + 1))
    circle_tile[:, :, 0] = 1
    origin_x, origin_y = np.size(circle_tile, 0)/2, np.size(circle_tile, 1)/2
    get_dist = lambda i, j: math.sqrt((i + 0.5 - origin_x)**2 + (j + 0.5 - origin_y)**2)

    for i in range(0, np.size(circle_tile, 0)):
        for j in range(0, np.size(circle_tile, 1)):

            r = get_dist(i, j)

            if r >= radius:
                f_B = 0
            else:
                f_B = 1

            f_A = 1 - f_B
            maj_material_ID = getMajorityMaterialID([f_A, f_B])
            circle_tile[i, j, :] = [f_A, f_B, maj_material_ID]

    return circle_tile

def populateBox(morphology):

    box = morphology.box
    coords = morphology.coords
    circle_tile = generateCircleTile(morphology.cfg)
    morphology.circle_tile = circle_tile

    def insertCircleTile(circle_tile, center, z):
        width = np.size(circle_tile, 0)
        minx = int(round(center[0] - width/2))
        miny = int(round(center[1] - width/2))
        maxx = int(round(minx + width))
        maxy = int(round(miny + width))

        # Check if cylinder is out of simulation bounds
        if minx >=0 and miny >=0 and maxx < len(box) and maxy < len(box):
            for i in range(maxx-minx):
                for j in range(maxy-miny):
                    if circle_tile[i, j, 0] != 1:
                        box[i+minx, j+miny, z, :] = circle_tile[i, j, :]

        return box

    for coord in coords:
        box = insertCircleTile(circle_tile, coord, 0)

    # We created the morphology as [X, Y, Z] but Cy-RSoXS takes [Z, Y, X]
    # Swap axes of our box
    box = np.swapaxes(box, 0, 2)
    return box

def generateMaterialMaps(box, num_materials):

    dimZ = np.size(box, 0)
    dimY = np.size(box, 1)
    dimX = np.size(box, 2)

    mat1_unaligned = box[:, :, :, 0]
    mat1_aligned = np.zeros((dimZ, dimY, dimX, 3))
    mat2_unaligned = box[:, :, :, 1]
    mat2_aligned = np.zeros((dimZ, dimY, dimX, 3))

    # The material maps define the fractional composition of material x in a given voxel
    # The element-wise sum of these arrays should give an array of all 1's
    # Currently, this check sums the unaligned maps because we do not support aligned fractions yet
    unaligned_sum = mat1_unaligned + mat2_unaligned
    unaligned_sum = unaligned_sum.flatten()

    if np.all(abs(unaligned_sum - 1) <= 1e-6) == False:
        raise ValueError('Uh oh! Not every voxel sums to 1')

    return [mat1_unaligned, mat2_unaligned], [mat1_aligned, mat2_aligned]

def generateHDF5(morphology, filepath: str = None, filename: str = None):

    f = h5py.File(filepath + filename + '.h5', 'w')

    g0   = f.create_group('Morphology_Parameters')
    dset_1 = g0.create_dataset('PhysSize', data = morphology.cfg.phys_size, dtype = np.float64)
    dset_2 = g0.create_dataset('NumMaterial', data = morphology.cfg.num_materials, dtype = np.float64)

    g1 = f.create_group('Vector_Morphology')
    dset_m1u = g1.create_dataset('Mat_1_unaligned', data=(morphology.unaligned[0]), dtype = np.float32, compression='gzip', compression_opts=9)
    g1['Mat_1_unaligned'].dims[0].label = 'Z'
    g1['Mat_1_unaligned'].dims[1].label = 'Y'
    g1['Mat_1_unaligned'].dims[2].label = 'X'
    dset_m2u = g1.create_dataset('Mat_2_unaligned', data=(morphology.unaligned[1]), dtype = np.float32, compression='gzip', compression_opts=9)
    g1['Mat_2_unaligned'].dims[0].label = 'Z'
    g1['Mat_2_unaligned'].dims[1].label = 'Y'
    g1['Mat_2_unaligned'].dims[2].label = 'X'

    dset_m1a = g1.create_dataset("Mat_1_alignment", data = (morphology.aligned[0]), dtype = np.float32, compression='gzip', compression_opts=9)
    g1['Mat_1_alignment'].dims[0].label = 'Z'
    g1['Mat_1_alignment'].dims[1].label = 'Y'
    g1['Mat_1_alignment'].dims[2].label = 'X'
    dset_m2a = g1.create_dataset("Mat_2_alignment", data = (morphology.aligned[1]), dtype = np.float32, compression='gzip', compression_opts=9)
    g1['Mat_2_alignment'].dims[0].label = 'Z'
    g1['Mat_2_alignment'].dims[1].label = 'Y'
    g1['Mat_2_alignment'].dims[2].label = 'X'

    f.close()

    
def generate_circle_lattice_test(return_folderpath=False, slurm_file=0, **slurm_kwargs):
    # default input parameters
    size_xy = 1024
    size_z = 1
    PhysSize = 2
    a_nm = 60
    b_nm = 60
    theta = 120
    radius_nm = 20
    NumMat = 2
    
    config = Config(size_xy = size_xy,
                size_z = size_z,
                phys_size = PhysSize,
                a_nm = a_nm,
                b_nm = a_nm,
                theta = theta,
                radius_nm = radius_nm)

    morphology = Morphology(config = config)
    morphology.makeMorphology()
    
    filepath = os.getcwd() + '/'
    filename = 'circle-lattice-test'
    generateHDF5(morphology, filepath=filepath, filename=filename)


    # write config and material files
    startEnergy = 285.0
    endEnergy = 286.0
    incrementEnergy = 1.0

    numThreads = 4;        #number of threads for execution
    #Files corresponding to Each material.
    dict={'Material0':'PEOlig2018.txt',
          'Material1':'vacuum'}


    # Label of energy to look for
    labelEnergy={"BetaPara":0,
                 "BetaPerp":1,
                 "DeltaPara":2,
                 "DeltaPerp":3,
                 "Energy":6}


    write_config2D(size_xy,PhysSize,[startEnergy,endEnergy],euler=False)

    # renamed "main" in generateConstants.py to something more descriptive
    write_materials(startEnergy,endEnergy,incrementEnergy,dict,labelEnergy,len(dict))

    # Create folder stamped with current datetime
    time_now = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    folder_name = 'circle_lattice_test_' + time_now
    folder_path = filepath + folder_name + '/'
    os.mkdir(folder_name)
    filename = filename + '.h5'
    subprocess.run(["mv","Material1.txt",f"./{folder_name}"])
    subprocess.run(["mv","Material2.txt",f"./{folder_name}"])
    subprocess.run(["mv","config.txt",f"./{folder_name}"])
    subprocess.run(["mv",f"{filename}",f"./{folder_name}"])

    if slurm_file == 1:
        write_slurm(filename, **slurm_kwargs)
        subprocess.run(["mv","job.slurm",f"./{folder_name}"])

    if return_folderpath:
        return folder_path
    
def run_slurm(folder_path):
    os.chdir(folder_path)
    subprocess.run(['sbatch','job.slurm'])
    

def test_circle_lattice(folder_path, savepng=False):
    h5_path = pathlib.Path(folder_path,'HDF5')
    h5s = sorted(list(h5_path.glob('*h5')))
    data = read_img_h5(h5s[0], PhysSize=2)
    remeshed_data = remesh(data)
    
    I_simulation = remeshed_data.mean('chi')
    
    # reference peak positions
    d_ref = 51.9615243
    q_star = 2*np.pi/d_ref
    n = np.array([1,3,4,7,9,12,13]) # first 7 peak orders
    q_ref = q_star*np.sqrt(n.astype(float))
    q_fit = []
    opt_params = []
    for peak in q_ref:
        q_lower = peak*0.9
        q_upper = peak*1.1
        I_fit = I_simulation.where((I_simulation.q > q_lower) & (I_simulation.q < q_upper) & ~np.isnan(I_simulation.values),drop=True)
#         print(I_fit.values)
        popt, pcov = curve_fit(gauss_bck, I_fit.q, I_fit.values,p0=(1e-2,peak,0.005), bounds=((0,q_lower,0),(10,q_upper,1)))
        q_fit.append(popt[1])
        opt_params.append(popt)
    
    q_fit = np.array(q_fit)
    
    assert(np.allclose(np.round(q_ref,3), np.round(q_fit,3)))
    
    

    if savepng:
        peaks_str = [r'q$^*$', r'$\sqrt{3}q^*$', r'$\sqrt{4}q^*$', 
                     r'$\sqrt{7}q^*$', r'$\sqrt{9}q^*$', 
                     r'$\sqrt{12}q^*$', r'$\sqrt{13}q^*$']
        centering = ['center','center','center','center','center','right','left']
        v_centering = ['center','center','center','center','center','top','bottom']
        with h5py.File(pathlib.Path(folder_path,'circle-lattice-test.h5'),'r') as f:
            model = f['Vector_Morphology']['Mat_2_unaligned'][()]
    
        plt.figure(constrained_layout=True,figsize=(3.25,2.7))
        plt.imshow(model[0,:,:],origin='lower',extent=[0,model.shape[1]*2,0,model.shape[2]*2])
        plt.xlabel('X [nm]')
        plt.ylabel('Y [nm]')
        plt.xlim(0,1000)
        plt.ylim(0,1000)
        plt.savefig(pathlib.Path(folder_path,'Model.png'),format='png',dpi=300,bbox_inches='tight')
        plt.close()
    

        plt.figure(constrained_layout=True,figsize=(3.25,2.7))
        I_simulation.plot(yscale='log')
        for i,peak in enumerate(q_ref):
            q_lower = peak*0.9
            q_upper = peak*1.1
            xnew = np.linspace(q_lower,q_upper)
            params = opt_params[i]
#             print(params,type(params))
            plt.plot(xnew,gauss_bck(xnew,*params)+1e-8,'k--')
            plt.text(peak,gauss_bck(peak,*params)*1.5,peaks_str[i],ha=centering[i],fontsize=6)
#             plt.plot(xnew,bck(xnew,*opt_params[i][-2:]),'k',linestyle='--')
        plt.xlabel(r'q [nm$^{-1}$]')
        plt.ylabel('I(q)')
        plt.xlim(4e-2,0.5)
        plt.ylim(1e-8)
        plt.savefig(pathlib.Path(folder_path,'PeakComparison.png'),format='png',dpi=300,bbox_inches='tight')
        plt.close()

        plt.figure(figsize=(3.25,2.7),constrained_layout=True)
        remeshed_data.plot(norm=LogNorm(1e-6,1e0))
        plt.title('Remeshed CyRSoXS Output')
        plt.xlabel(r'q [nm$^{-1}$]')
        plt.ylabel(r'$\chi$ [Degrees]')
        plt.savefig(pathlib.Path(folder_path,'remeshed.png'),format='png',dpi=300,bbox_inches='tight')
        plt.close()

        plt.figure(figsize=(3.25,2.7),constrained_layout=True)
        data.plot(norm=LogNorm(1e-6,1e0),cmap='jet')
#         plt.title('2D CyRSoXS Output')
        plt.xlabel(r'q$_x$ [nm$^{-1}$]')
        plt.ylabel(r'q$_y$ [nm$^{-1}$]')
        plt.savefig(pathlib.Path(folder_path,'raw.png'),format='png',dpi=300,bbox_inches='tight')
        plt.close()
        
        plt.figure(figsize=(3.25,2.7),constrained_layout=True)
        data.plot(norm=LogNorm(1e-6,1e0),cmap='jet')
        plt.xlim(0,0.5)
        plt.ylim(0,0.5)
#         plt.title('2D CyRSoXS Output')
        plt.xlabel(r'q$_x$ [nm$^{-1}$]')
        plt.ylabel(r'q$_y$ [nm$^{-1}$]')
        for i,peak in enumerate(q_ref):
            plt.text(0.03,peak,peaks_str[i],va=v_centering[i],ha='center',fontsize=6)
        plt.savefig(pathlib.Path(folder_path,'raw_zoomed.png'),format='png',dpi=300,bbox_inches='tight')
        plt.close()
        
    
   
    
    
def gauss_bck(x,a,b,c):
    return a*np.exp(-(x-b)**2/c**2)

def bck(x,d,e):
    return d*x + e
    
    
    
    
    
if __name__ == '__main__':
    version = '1.0.1.1'
    slurm_file = 0
    
    for arg in sys.argv:
        if arg.startswith('slurm'):
            slurm_file = int(arg.split('=')[-1])
        if arg.startswith('version'):
            version = arg.split('=')[-1]

    folderpath = generate_circle_lattice_test(return_folderpath=True,slurm_file=slurm_file,cyrsoxs_version=version)
    
    if slurm_file == 1:
        import time
        run_slurm(folderpath)
        hdf5path = pathlib.Path(folderpath,'HDF5')
        while len(list(hdf5path.glob('*h5'))) < 1:
            time.sleep(1)
        test_circle_lattice(folderpath,savepng=True)


