from refl1d.names import *
from refl1d.probe import make_probe 

from copy import copy
import numpy as np
import sys
import os
from types import MethodType
from numpy import inf
from pathlib import Path
import matplotlib.pyplot as plt

from refl1d.names import *
from copy import copy
import numpy as np
import sys
import os
from types import MethodType
from numpy import inf
from refl1d.probe import make_probe 
from bumps.parameter import Parameter 
from bumps.util import push_seed
from pathlib import Path
import matplotlib.pyplot as plt

from Load_data import rhoM_Calc


#### Multilayer Objects ####

"""

Standard slab based model with coupled structural parameters
For magnetic fitting paramaters & models can be passed through the multilayer.magnetise method

"""

# keep updated with the name of any magnetic materials used:



def Assemble_stack(obj, mml, full_struct, give_magnetism = ['CoB','CoFeB','Fe','Co','Py','Ni','Pt', 'PIM'],
                    thickness_gradients = {}, roughness_gradients = {}, magnetic_modelling = True,
                    **kwargs):

    """
    obj : object to assign created sample to, can be different to mml to carry over fitted params between models
    mml : existing multilayer object containing relavent sample property dictionaries
    full_struct: list of refl1D materials in mutliayer stack: full_struct = Buffer + Stack*repeats + Cap
    gradients : dictionaries with keys of material names, otherwise set to 0 if missing values
    give_magnetism : list of stack names to asign magnetism, example set given 
    """

    # store passed values

    obj.thickness_gradients = thickness_gradients ; obj.roughness_gradients = roughness_gradients

    for i, x in enumerate(full_struct):

        density = x.density
        obj.Densities[x.name] = density 

        # handle option for structural gradients

        if x.name in thickness_gradients:
            thickness_gradient = thickness_gradients[x.name]
        else:
            thickness_gradient = 0

        if x.name in roughness_gradients:
            roughness_gradient = roughness_gradients[x.name] 
        else:
            roughness_gradient = 0

        # assembles the slab model

        if i == 0:
            # substrate layer is defined first
            obj.sample = (Slab(material = x, thickness = mml.Thicknesses[x.name], interface = mml.Roughnesses[x.name],
                             ))

        elif x.name in give_magnetism and magnetic_modelling:

            # allows layer to be magnetised later by "magnetise" method

            obj.sample =  obj.sample | Slab(material = x, thickness = mml.Thicknesses[x.name] + i * thickness_gradient,
             interface = mml.Roughnesses[x.name] + i * roughness_gradient,
                             magnetism=Magnetism(rhoM=None,  interface_below = None,
                                                 interface_above = None))

        else:
            # default non-magnetic layer
            obj.sample =  obj.sample | Slab(material = x, thickness = mml.Thicknesses[x.name] + i * thickness_gradient,
             interface = mml.Roughnesses[x.name] + i * roughness_gradient)
            
    
    obj.sample = obj.sample | air

    return(obj)


class Multilayer:
    
    
    def __init__(self, Buffer : list, Stack : list, Cap : list,
                 thicknesses, roughnesses = 1, repeats = 6, **kwargs):
        """
        Parameters
        ----------
        Buffer : list
            DESCRIPTION.
            List of Buffer Refl1D materials objects
        Stack : list
            DESCRIPTION.
            List of main stack Refl1D materials objects
        Cap : list
            DESCRIPTION.
            List of Cap Refl1D materials objects
        thicknesses : TYPE
            DESCRIPTION.
            Thicknesses of layers, give as a list in order [Buffer, Stack, Cap] or as a dictionary as {name : value}
        roughnesses : TYPE, optional
            DESCRIPTION. The default is 1 Angstrom.
            Roughness of layers, give as a constant, a list in order [Buffer, Stack, Cap] or as a dictionary as {name : value}
        repeats : TYPE, optional
            DESCRIPTION. The default is 6.
            Number of main stack repeats
        """

        full_struct = Buffer + Stack*repeats + Cap 

        self.struct = [x.name for x in full_struct]

        #store initial configs
        self.Buffer = Buffer ; self.Stack = Stack ; self.Cap = Cap; self.init_thicknesses = thicknesses; self.init_roughnesses = roughnesses; self.repeats = repeats 
                      
        self.Roughnesses = self.assign(roughnesses)
        self.Thicknesses = self.assign(thicknesses)
        self.RhoMs = self.assign(0)
        self.Densities = {}
    
        Assemble_stack(self, self, full_struct, **kwargs)
    
    def assign(self, value):
        
        if type(value) in (float, int):
            # assign globally, e.g. constant roughness
            values = {}
            for x in (self.struct):
                values[x] = value
            return(values)
        
        elif type(value) is list:
            # assign values to layers in ascending order
            values = {}
            for i, x in enumerate(self.Buffer + self.Stack + self.Cap):
                values[x.name] = value[i] 
            
            return(values)
        
        elif type(value) is dict:
            # assign to specified layers
            values = value
            return(values)
        
        else:
            raise TypeError('Check Parameter Class in Assignment')

    def __str__(self):

        # lists a summary of multilayer values 

        print('Index  |  Composition   |  Thickness  | Roughness  | RhoM')
        
        print('----------------------------------------------------------')

        for i, layer in enumerate(self.sample):

            if layer.name != 'air':

                try:
                    x = layer.thickness.value

                    # trys to display with .value (if stack data is bumps fitting parameter type)

                    print(f"    {str(i)}  |  {str(layer.name)}      |     {str(layer.thickness.value)}       |       {str(layer.interface.value)}      |     {str(layer.magnetism.rhoM.value)}")
                
                except AttributeError:

                    print(f"    {str(i)}  |  {str(layer.name)}      |     {str(layer.thickness)}       |       {str(layer.interface)}      |     {str(layer.magnetism.rhoM)}")
    
    
    def magnetise(self, layer_name, rhoM, layer_num = None,
                  interface_below = None, interface_above = None,
                  dead_above = None, dead_below = None, rhoM_gradient = 0,
                  rhoM_roughness_gradient = 0, **kwargs):
        """
        Magnetises a specified layer

        layer_name :
            Specify the name of the layer to give magnetism
        layer_num: 
            specify a perticular layer to magnetise, if None, magnetises all layers of specified name
        moment:
            rhoM magnitude or fitting parameter 
        interfaces:
            for Fitting parameter. if None, magnetic interface will share roughness of the structural model
        dead layers:
            for Fitting parameter
        gradient:
            applies a gradient over the 
        """
        counter = 0

        for i, layer in enumerate(self.sample):

            if layer.name == layer_name:

                counter += 1

                if layer_num is None or layer_num == counter:
                    
                    self.sample[i].magnetism.rhoM = (rhoM + i * rhoM_gradient)


                    if interface_below is not None:
                        self.sample[i].magnetism.interface_below = (interface_below + i * rhoM_roughness_gradient) 
                    if interface_above is not None:
                        self.sample[i].magnetism.interface_above = (interface_above + i * rhoM_roughness_gradient) 

                    if dead_above is not None:
                        self.sample[i].magnetism.dead_above = (dead_above)
                    if dead_below is not None:
                        self.sample[i].magnetism.dead_below = (dead_below)

                    # store

                    self.RhoMs[layer_name] = rhoM
        
    def plot_reflectivity(self, probe, fig = None):
        
        # calls built in refl1d function of same name for provided probe

        M = Experiment(probe = probe, sample = self.sample)
        M.simulate_data(noise= 2)
        
        if fig is None:
            fig = plt.figure(figsize = (6,4))
        else:
            plt.figure(fig.number)
            
        M.plot_reflectivity()
        plt.tight_layout()
    
        return(fig)
    
    def plot(self, probe, fig = None):

        # calls built in refl1d function of same name for provided probe
        
        M = Experiment(probe = probe, sample = self.sample)
        M.simulate_data(noise= 2)
        
        if fig is None:
            fig = plt.figure(figsize = (6,4))
        else:
            plt.figure(fig.number)
            
        M.plot()
        
        plt.tight_layout()
        
        return(fig)

    def duplicate_structure(self, **kwargs):

    # duplicates the mutilayer object with shared fitting parameters but allowing for different magnetic properties - used in multi-field fittings

        copy = Copy_Multilayer(self, **kwargs)

        return(copy)



class FitMultilayer(Multilayer):

    def set_structural(self, roughnesses_ranges = {}, thickness_percentage_var = 40,
                       ignore_layers = [], obj = None, default_roughness_range = (1,3),
                       **kwargs):
            
            # turns initial value estimates into fitting parameters
            # assumes 40% tolerence by default in thicknesses and a 1-3 A roughness range, for rougher layers specify in a "roughness_ranges" dict

            if obj is None:
                # a different obj allows asigning to a separate object without overiding float params
                obj = self

            for i, layer in enumerate(self.struct):
                
                if layer not in ignore_layers:
                    
                    name = layer + " roughness"
                    
                    if layer in roughnesses_ranges:
                        
                        roughness_range = roughnesses_ranges[layer]

                    else:

                        roughness_range = default_roughness_range
                        
                    obj.Roughnesses[layer] = Parameter.default(self.Roughnesses[layer], name = name)
                    obj.Roughnesses[layer].range(roughness_range[0], roughness_range[1])
                    
                    name = layer + " thickness"
                    
                    if i != 0:

                        obj.Thicknesses[layer] = Parameter.default(self.Thicknesses[layer], name = name)
                        obj.Thicknesses[layer].pmp(-thickness_percentage_var,thickness_percentage_var) # a 40% default error
                        
                    else:

                        # substrate thickness should be zero and NOT fitted

                        obj.Thicknesses[layer] = 0


                    if layer[-4:] == 'cap' and False:

                        # couple cap density to fitting parameter of main value, not currently working?

                        like_material = layer[:-4]
                        obj.Densities[layer] = obj.density_parameters[like_material]

                    
                    else:

                        # fit density

                        name = layer + " density"
                          
                        obj.Densities[layer] = Parameter.default(self.Densities[layer], name = name)
                        obj.Densities[layer].pmp(-20, 20)
                    
                        
                for slab in obj.sample:
                    
                    if slab.name == layer:

                        # store values, allows carry over of parameters by .duplication method

                        slab.thickness = self.Thicknesses[layer]
                        slab.roughness = self.Roughnesses[layer]
                        slab.density = self.Densities[layer]


        
    def __init__(self, Buffer : list, Stack : list, Cap : list, thicknesses: list,
                gradient = False,
                 **kwargs):

        # make simulated multilayer with initial values
        super().__init__(Buffer, Stack, Cap, thicknesses,  **kwargs)

        self.thickness_gradients = {}
        self.roughness_gradients = {}

        if gradient:

            for x in Stack:
                
                self.thickness_gradients[x.name] = Parameter.default(0, name= f'{x.name} Thickness gradient')
                self.thickness_gradients[x.name].range(-1, 1)
                self.roughness_gradients[x.name] = Parameter.default(0, name= f'{x.name} Roughness gradient')
                self.roughness_gradients[x.name].range(-1, 1)

        # turn initial values from floats to parameters
        self.set_structural(**kwargs) 

        # remove roughnesses as kwarg so can be set as fitted var
        kwargs = {key: value for key, value in kwargs.items() if key != 'roughnesses'}

        # reinitialise with parameters as stack values
        super().__init__(Buffer, Stack, Cap, self.Thicknesses, roughnesses = self.Roughnesses, thickness_gradients = self.thickness_gradients,
            roughness_gradients = self.roughness_gradients, **kwargs)
    

class Copy_Multilayer(FitMultilayer):

    def __init__(self, copy, **kwargs):

        # copies the .sample properties of another multilayer, carrying over fitting params

        full_struct = copy.Buffer + copy.Stack*copy.repeats + copy.Cap

        self.Buffer = copy.Buffer ; self.Stack = copy.Stack ; self.Cap = copy.Cap; self.repeats = copy.repeats 

        self.Densities = {}
        self.struct = [x.name for x in full_struct]

        Assemble_stack(self, copy, full_struct, thickness_gradients = copy.thickness_gradients,
            roughness_gradients = copy.roughness_gradients, **kwargs)

        self.RhoMs = {}      # allows you to call .magnetise
        self.Roughnesses = self.assign(copy.Roughnesses)
        self.Thicknesses = self.assign(copy.Thicknesses)

