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

#### Magnetic Fitting Models ####


# hard axis anisotropy functions

# 'a' taken as a phenomelogical stiffness parameter 


coth = lambda x : np.cosh(x) / np.sinh(x)   

def coth_inf(x):
    y = coth(x)
    for i, v in enumerate(y):
        if np.isnan(v):
            y[i] = 1
    return(np.array(y))
    
Langevin = lambda x, a : coth(x*a) - 1/(x*a)

Brilluin = lambda x, a : np.tanh(x/a)


def Assign_uniform_magnetisation_values(multilayer, fields : list,
    rhoM_roughness = None,  **kwargs):

    """
    Uniformly assigns all magnetic layers in a stack a fitting parameter for each field
    Magnetic roughness parameter can be passed in kwargs
    """

    samples = []
    
    for i, field in enumerate(fields):

        name = f'rhoM_field_{field}'


        rhoM = Parameter.default(rhoM_Calc(600), name= name)

        rhoM.range(rhoM_Calc(0), rhoM_Calc(900))
        
        sample = multilayer.duplicate_structure()
        
        sample.magnetise('CoB', rhoM, interface_below = rhoM_roughness,
         interface_above = rhoM_roughness,  **kwargs)
        
        samples += [sample]
           
    return(samples)
       


def Assign_uniform_magnetisation_values_PIM(multilayer, fields : list,
    rhoM_roughness = None,  PIM_roughness = None,**kwargs):

    """
    Uniformly assigns all magnetic layers in a stack a fitting parameter for each field
    Magnetic roughness parameter can be passed in kwargs
    """

    samples = []
    
    for i, field in enumerate(fields):

        name = f'rhoM_field_{field}'


        rhoM = Parameter.default(rhoM_Calc(600), name= name)

        rhoM.range(rhoM_Calc(0), rhoM_Calc(900))
        
        sample = multilayer.duplicate_structure()
        
        sample.magnetise('CoB', rhoM, interface_below = rhoM_roughness,
         interface_above = rhoM_roughness,  **kwargs)


        name = f'PIM_field_{field}'


        PIM = -Parameter.default(rhoM_Calc(600), name= name)

        PIM.range(-rhoM_Calc(900), rhoM_Calc(900))

        sample.magnetise('Pt', PIM, interface_below = PIM_roughness,
         interface_above = PIM_roughness,  **kwargs)
        
        samples += [sample]
           
    return(samples)

def Assign_uniform_magnetisation_function(multilayer, fields : list, *args,
    rhoM_roughness = None, function = Langevin,  **kwargs):

    """

    Models rhoM as a function of field uniformly throughout the stack
    Magnetic roughness parameter can be passed in kwargs

    rhoM - Saturated magnetisation IP, can be fitting param
    function - Function relating rhoM to "fields"
    *args - function params - can be fitting params
    rhoM_roughness - can be fitting param, for separate magnetic interfaces
    """
    n = multilayer.repeats  
    samples = []
    
    rhoMs = {}
    func_params = {}

    for i in range(1,n+1):
            
        name = f'IP_stiffness_layer_{i}'

        func_param = Parameter.default(20, name= name)
        func_param.range(3,3000)

        func_params[str(i)] = func_param

    for f in fields:

        sample = multilayer.duplicate_structure()

        for i in range(1,n+1):  

            rhoM_i = function(f, func_params[str(i)])
            sample.magnetise('CoB',  rhoM_i, layer_num = i, interface_below = rhoM_roughness,
             interface_above = rhoM_roughness,  **kwargs)
            
        samples += [sample]
               
    return(samples)


def Assign_uniform_magnetisation_function(multilayer, fields : list, *args,
    rhoM_roughness = None, function = Langevin, oxidation = False, **kwargs):

    """

    Models rhoM as a function of field uniformly throughout the stack
    Magnetic roughness parameter can be passed in kwargs

    rhoM - Saturated magnetisation IP, can be fitting param
    function - Function relating rhoM to "fields"
    *args - function params - can be fitting params
    rhoM_roughness - can be fitting param, for separate magnetic interfaces
    """
    n = multilayer.repeats  
    samples = []
    
    rhoMs = {}
    func_params = {}

    Ms = Parameter.default(2, name= 'Ms')
    Ms.range(1,5)

    if oxidation:

        oxidation_percentage = Parameter.default(1, name= 'Oxidation percentage')
        oxidation_percentage.range(0.1,1)

    for i in range(1,n+1):
            
        name = f'IP_stiffness_layer_{i}'

        func_param = Parameter.default(20, name= name)
        func_param.range(3,200)

        func_params[str(i)] = func_param

    for f in fields:

        sample = multilayer.duplicate_structure()

        for i in range(1,n+1):  

            rhoM_i = Ms*function(f, func_params[str(i)])

            if oxidation and i == n:
                    
                rhoM_i *= oxidation_percentage

            sample.magnetise('CoB',  rhoM_i, layer_num = i, interface_below = rhoM_roughness,
             interface_above = rhoM_roughness,  **kwargs)
            
        samples += [sample]
               
    return(samples)


def Assign_nonuniform_magnetisation_function(multilayer, fields : list, *args,
    rhoM_roughness = None, function = Langevin,  **kwargs):

    """

    Models rhoM as a function of field uniformly throughout the stack
    Magnetic roughness parameter can be passed in kwargs

    rhoM - Saturated magnetisation IP, can be fitting param
    function - Function relating rhoM to "fields"
    *args - function params - can be fitting params
    rhoM_roughness - can be fitting param, for separate magnetic interfaces
    """
    n = multilayer.repeats  
    samples = []
    Ms_s = n*[0,]
    rhoMs = {}
    func_params = {}

    for i in range(1,n+1):
            
        name = f'IP_stiffness_layer_{i}'

        func_param = Parameter.default(20, name= name)
        func_param.range(3,3000)

        func_params[str(i)] = func_param

        name = f'Ms_{i}'

        Ms = Parameter.default(rhoM_Calc(700), name= name)
        Ms.range(rhoM_Calc(400),rhoM_Calc(900))

        Ms_s[i-1] = Ms

    for f in fields:

        sample = multilayer.duplicate_structure()

        for i in range(1,n+1):  

            rhoM_i = Ms_s[i-1]*function(f, func_params[str(i)])
            sample.magnetise('CoB',  rhoM_i, layer_num = i, interface_below = rhoM_roughness,
             interface_above = rhoM_roughness,  **kwargs)
            
        samples += [sample]
               
    return(samples)


def Assign_gradient_magnetisation_values(multilayer, fields : list, decouple = [],
     **kwargs):

    """
    More complex gradient model allowing for decoupled layers and gradients both in rhoM and rhoM's roughness
    """

    rhoM_roughness = Parameter.default(1, name= 'RhoM Roughness')
    rhoM_roughness.range(0.1, 6)

    # roughness gradient used across uncoupled layers

    rhoM_roughness_gradient = Parameter.default(-0.2, name= 'RhoM Roughness Gradient')
    rhoM_roughness_gradient.range(-1, 0)

    rhoM_roughness = Parameter.default(1, name= 'RhoM Roughness')
    rhoM_roughness.range(0.1, 6)

    rhoM_gradient = Parameter.default(0, name= 'RhoM Gradient')
    rhoM_gradient.range(-rhoM_Calc(100), rhoM_Calc(100))


    samples = []
    
    n = multilayer.repeats

    for j, field in enumerate(fields):

        name = f'rhoM_field_{field}'


        rhoM_uni = Parameter.default(rhoM_Calc(600), name= name)

        rhoM_uni.range(-rhoM_Calc(900), rhoM_Calc(900))
        
        sample = multilayer.duplicate_structure()
        

        for i in range(1, n+1):

            if i in decouple:

                name = f'rhoM_layer{i}_field_{field}'

                rhoM = Parameter.default(rhoM_Calc(600), name= name)

                rhoM.range(-rhoM_Calc(900), rhoM_Calc(900))

                sample.magnetise('CoB', rhoM, interface_below = rhoM_roughness,
                    interface_above = rhoM_roughness, layer_num = i, rhoM_gradient =rhoM_gradient,
                     rhoM_roughness_gradient = rhoM_roughness_gradient,  **kwargs)

            else:

                sample.magnetise('CoB', rhoM_uni, interface_below = rhoM_roughness,
                    interface_above = rhoM_roughness,  layer_num = i, rhoM_gradient =rhoM_gradient,
                     rhoM_roughness_gradient = rhoM_roughness_gradient, **kwargs)

            
        samples += [sample]
           
    return(samples) 


def Proximity_Magnetise(samples, **kwargs):
    PIM_samples = []

    for i, sample in enumerate(samples):

        name = f'PIM_{i}_rhoM'

        rhoM = Parameter.default(0, name= name)
        rhoM.range(-rhoM_Calc(900), rhoM_Calc(900))

        name = f'PIM_{i}_roughness'

        rhoM_roughness = Parameter.default(5, name= name)
        rhoM_roughness.range(3, 30)

        sample.magnetise('PIM', rhoM, interface_below = rhoM_roughness,
                        interface_above = rhoM_roughness,  layer_num = i, **kwargs)

        PIM_samples += [sample]

    return(PIM_samples)


def Assign_magnetisation_values_decoupled_layers(multilayer, fields : list, decouple = [],
    rhoM_roughness = None,  **kwargs):

    """
    Uniformly assigns all magnetic layers in a stack a fitting parameter for each field
    If layers are specified in decouple, will be modeled seperatley
    Magnetic roughness parameter can be passed in kwargs
    """

    samples = []
    n = multilayer.repeats  

    for j, field in enumerate(fields):

        if type(rhoM_roughness) is list:

            rhoM_R = rhoM_roughness[j]

        else:

            rhoM_R = rhoM_roughness


        name = f'rhoM_field_{field}'


        rhoM_uni = Parameter.default(rhoM_Calc(600), name= name)

        rhoM_uni.range(rhoM_Calc(0), rhoM_Calc(1200))
        
        sample = multilayer.duplicate_structure()
        

        for i in range(1, n+1):

            if i in decouple:

                name = f'rhoM_layer{i}_field_{field}'

                rhoM = Parameter.default(rhoM_Calc(600), name= name)

                rhoM.range(rhoM_Calc(0), rhoM_Calc(900))

                sample.magnetise('CoB', rhoM, interface_below = rhoM_R,
                    interface_above = rhoM_R, layer_num = i,  **kwargs)

            else:

                sample.magnetise('CoB', rhoM_uni, interface_below = rhoM_R,
                    interface_above = rhoM_R,  layer_num = i, **kwargs)

            
        samples += [sample]
           
    return(samples)


def Assign_magnetisation_values_decoupled_rough_layers(multilayer, fields : list, decouple = [],
    rhoM_roughness = None, gradient = True,  **kwargs):

    """
    Uniformly assigns all magnetic layers in a stack a fitting parameter for each field, assigns separate magnetic roughnesses
    """

    samples = []
    n = multilayer.repeats  

    # make all magnetic roughness variables (not field dependent)

    rhoM_roughness = Parameter.default(1, name= 'RhoM Roughness')
    rhoM_roughness.range(0.1, 6)


    rhoM_roughnesses = {}

    # roughness gradient used across uncoupled layers

    if gradient:

        rhoM_roughness_gradient = Parameter.default(-0.2, name= 'RhoM Roughness Gradient')
        rhoM_roughness_gradient.range(-1, 0)

    else:

        rhoM_roughness_gradient = 0

    for i in range(1, n+1):

        if i in decouple:

            rhoM_roughnesses[str(i)]  = Parameter.default(1, name= f'RhoM Roughness layer {i}')
            rhoM_roughnesses[str(i)].range(0.1, 6)


    for j, field in enumerate(fields):

        name = f'rhoM_field_{field}'


        rhoM_uni = Parameter.default(rhoM_Calc(600), name= name)

        rhoM_uni.range(rhoM_Calc(0), rhoM_Calc(900))
        
        sample = multilayer.duplicate_structure()
        

        for i in range(1, n+1):

            if i in decouple:

                name = f'rhoM_layer{i}_field_{field}'

                rhoM = Parameter.default(rhoM_Calc(600), name= name)

                rhoM.range(rhoM_Calc(0), rhoM_Calc(900))

                sample.magnetise('CoB', rhoM, interface_below = rhoM_roughnesses[str(i)],
                    interface_above = rhoM_roughnesses[str(i)], layer_num = i,  **kwargs)

            else:

                sample.magnetise('CoB', rhoM_uni, interface_below = rhoM_roughness,
                    interface_above = rhoM_roughness,  layer_num = i, **kwargs)

            
        samples += [sample]
           
    return(samples)


def Assign_magnetisation_convolving_function_per_layer(multilayer, fields : list,
    rhoM_roughness = None, function = Langevin,  **kwargs):

    """
    Convolves effect of varying anisotropy and increasing field

    rhoM - Saturated magnetisation IP
    function - Function relating rhoM to "fields"
    rhoM_roughness - can be fitting param, for separate magnetic interfaces
    """

    n = multilayer.repeats  
    samples = []
    
    rhoMs = {}
    func_params = {}

    for f in fields:

        name = f"Max_rhoM_field_{f}T"

        rhoM = Parameter.default(rhoM_Calc(600), name= name)

        rhoM.range(rhoM_Calc(0), rhoM_Calc(900))

        rhoMs[str(f)] = rhoM

    for i in range(1,n+1):
            
        name = f'IP_stiffness_layer_{i}'

        func_param = Parameter.default(20, name= name)
        func_param.range(3,3000)

        func_params[str(i)] = func_param

    for f in fields:

        sample = multilayer.duplicate_structure()

        for i in range(1,n+1):  

            rhoM_i = rhoMs[str(f)] * function(f, func_params[str(i)])
            sample.magnetise('CoB',  rhoM_i, layer_num = i, interface_below = rhoM_roughness,
             interface_above = rhoM_roughness,  **kwargs)
            
        samples += [sample]
               
    return(samples)


def Assign_magnetisation_per_layer_per_field(multilayer, fields, rhoM_roughness = None, **kwargs):

    # for when nothing else works...

    n = multilayer.repeats  
    samples = []
    
    for f in fields:

        sample = multilayer.duplicate_structure()

        for i in range(1,n+1):
            
            name = f'rhoM_layer_{i}_field_{f}'

            rhoM = Parameter.default(rhoM_Calc(600), name= name)

            rhoM.range(rhoM_Calc(0), rhoM_Calc(900))

            sample.magnetise('CoB',  rhoM, layer_num = i, interface_below = rhoM_roughness,
             interface_above = rhoM_roughness,  **kwargs)
            
        samples += [sample]
               
    return(samples)
  
