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


def logstep(start, stop, step, base=10.0):
    """
    Creates a log spaced 1d array by defining a step size and a base
    In the form of dQ/Q - i.e. dQ\\Q*Qpoint
    """
    logrange = [start]
    point = start
    while point < stop:
        point = point+base**(np.log10(step*point)/np.log10(base))
        # print(f"point = {point}")
        logrange.append(point)
    # print(f"Q={logrange}")
    return np.array(logrange)


def QT2L(Q,T):
    '''
    Simple Q,T to L conversion
    
    Q = 4 pi sin(T) / L   ==>  L = 4 pi sin(T) / Q
    '''
    return 4 * np.pi * np.sin(np.radians(T)) / Q


def rhoM_Calc(Mag,Units='emu_cc',N=None):
    '''
    Converts Magnetic units to mSLD. 
    N is the density per FU and is only required if converting from uB/FU
    '''
    if Units=='emu_cc':
        M = 1000*Mag
        rhoM = M * 2.853e-6
    elif Units=='kA/m':
        rhoM = Mag * 2.91043e-3 # From Kirby
    else:
        # in uB
        rhoM = 2.645e-5*Mag*N*1e6
        #rhoM = 9.274e-21*Mag
    return rhoM


def TOF_loader(T=0.25, dQoQ=0.02, 
               Q_sim_range=(0.005, 0.2),
               filename=None, skiprows=1, **kw):
    """
    Loads and creates NeutronProbe objects for TOF stitched datasets
    I.e. from multiple angles. In the case of ISIS NR instruments we 
    typically have a constant dq/q resolution which the data is binned to at
    the end of the reduction.

    *T* incident theta for the lowest angle
    *dQoQ* dq/q resolution data has been binned to
    *filename* filename of the data set to be loaded
    *kw* keyword arguments (kwargs) to be passed to NeutronProbe()
    """

    # np.loadtext is currently set for simple 3 column POLREF data
    # if you aim to use this loader for other data, talk to your local contact
    # to understand the data format, and how best to load it.

    if filename is not None:
        data = np.loadtxt(filename, skiprows=skiprows).T
        if dQoQ is None:
            Q, R, dR, dQo = data
        else:
            Q, R, dR = data
            dQo = (Q*dQoQ)
        data_in = (R, dR)
    else:
        Q = logstep(Q_sim_range[0], Q_sim_range[1], dQoQ, base=dQoQ)
        data_in = None

    L = QT2L(Q, T)
    # Converting the dq/q resolution into a dq value for each Q point
    # dQ = FWHM2sigma(dQo)
    # Since we take dL/L = 0, dQ/Q = dT/T, so dT = T * dQoQ
    dT = T *dQoQ

    probe_out = NeutronProbe(
        T=T, 
        dT=dT,
        L=L, 
        dL=0, 
        data=data_in,
        # For standard TOF measurements resolution is assumed to be normal (gaussian)
        #  For measurements with many wavelengths and many angles (say cw measurements)
        #  then a uniform resolution can be used instead.
        resolution='normal',
        **kw 
    )

    return probe_out


def load_probe_polref(filename, angle, dQoQ, name=None, path=None, pol_mode=None, field=None, **kw):
    """
    creates one probe (Neutron, Polarized - PA or PNR) from one measurement - could be one angle or stitched dataset.
    If polarized, sets some default values and links instrumental parameters for each cross-section together
    """

    if name is None:
        name = filename
    if path is None:
        path = os.getcwd()
    
    filepath = Path(path)/filename
    
    if (pol_mode != "pnr") and (pol_mode != "pa"):
        probe = TOF_loader(T=angle, dQoQ=dQoQ, filename=f"{filepath}.dat", name=name, **kw)

    else:
        if pol_mode == "pa":
            files = dict(data_mm=f"{filepath}_dd.dat",
                         data_mp=f"{filepath}_du.dat",
                         data_pm=f"{filepath}_ud.dat",
                         data_pp=f"{filepath}_uu.dat")
        else:
            files = dict(data_mm=f"{filepath}_d.dat",
                         data_mp=None,
                         data_pm=None,
                         data_pp=f"{filepath}_u.dat")

        cross_sections = []
        for data in files.values():
            if data is None:
                cross_sections.append(None)
            else:
                cross_sections.append(TOF_loader(T=angle, dQoQ=dQoQ, filename=data, name=name, **kw))
        if field is None:
            field = 0.0

        probe = PolarizedNeutronProbe(cross_sections, Aguide=270, H=field, name=name)

        for xs in (probe.mm, probe.mp, probe.pm, probe.pp):
            if xs is not None:
                xs.name = name
                xs.intensity = probe.pp.intensity
                xs.sample_broadening = probe.pp.sample_broadening
                xs.theta_offset = probe.pp.theta_offset
                xs.background = probe.pp.background

        probe.pp.intensity.name = f"intensity {name}"
        probe.pp.background.name = f"background {name}"
        probe.pp.sample_broadening.name = f"sample_broadening {name}"

        probe.pp.intensity.tags = ["inst", "nuisance"]
        probe.pp.background.tags = ["inst", "nuisance"]
        probe.pp.sample_broadening.tags = ["inst", "nuisance"]

    return probe


def load_X_ray_probe(filename, name = None, path = None, limit_T = None, norm = True,
             intensity = 1,background = 1e-5, sample_broadening = 0.0, theta_offset = 0):

    if name is None:
        name = filename
    if path is None:
        path = os.getcwd()


    T, R = np.loadtxt(filename, unpack = True)
    T = np.array(T)/2 ; R = np.array(R)  # converts 2 theta to theta

    if limit_T is not None and type(limit_T) is list:

        # snip 2theta range between min and max T

        R = R[T > limit_T[0]];  T = T[T > limit_T[0]]
        R = R[T < limit_T[1]] ; T = T[T < limit_T[1]]

    #dR = (R)*0.1 + min(R)  # force >10% errors in brute force way of keeping chi2 down
    dR = (R)*0.1 + R**0.5 + min(R)
    if norm:
        R0 = max(R)
        R = R/R0
    else:
        R0 = 1

    dR = dR/R0


    data = (R, dR)

    probe = XrayProbe(data = data , L = 1.54, dL = 0.005, T = T, dT = 0.001, intensity = intensity, background = background,
     sample_broadening = sample_broadening, theta_offset = theta_offset, name = name, resolution = 'normal')

    probe.intensity.name = f"intensity {name}"
    probe.background.name = f"background {name}"
    probe.sample_broadening.name = f"sample_broadening {name}"
    probe.theta_offset.name = f"theta offset {name}"
    probe.Aguide = Parameter.default(probe.Aguide, name= 'probe Aguide')  # not including this throws an error
    probe.H = Parameter.default(1, name= 'probe H') # not including this throws an error

    probe.intensity.tags = ["inst", "nuisance"]
    probe.background.tags = ["inst", "nuisance"]
    probe.sample_broadening.tags = ["inst", "nuisance"]
    probe.theta_offset.tags = ["inst", "nuisance"]
    probe.Aguide.tags = ["nuisance"]
    probe.H.tags = ["nuisance"]

    probe.intensity.range(1e-1, 20000000000000)
    probe.background.range(1e-9, 1e2)
    probe.sample_broadening.range(0, 0.03)
    probe.theta_offset.range(-0.05,0.05)

    return(probe)



### Data Loader ###


def probe_loader(code,
    dQoQ = 0.01, theta = 0.25):

    code = code.upper()  # case insensitive

    if code == "TI_MO_ML":
        path = r"/data/analysis/POLREF/RBNumber/RB2510555/"
        files = [r"TI_MO_ML_RT_68Amps", r"TI_MO_ML_RT_20Amps", r"TI_MO_ML_RT_10Amps", r"TI_MO_ML_RT_5Amps" , r"TI_MO_ML_RT_1Amp"]
        names = ["Mo(10 A), I = 68 A", "Mo(10 A), I = 20 A", "Mo(10 A), I = 10 A", "Mo(10 A), I = 5 A", "Mo(10 A), I = 1 A"]

    elif code == "TI_TA15_ML":
        path = r"/data/analysis/POLREF/RBNumber/RB2410236/"  
        files = [r"MET20240704-01-0.7T", r"MET20240704-01-0.2T", r"MET20240704-01-0.01T"]
        names = ["Ta(15 A), I = 68 A", "Ta15(15 A), I = 20 A", "Ta15 (15 A), I = 1 A"]

    elif code == 'TI_ML':
        path = r"/data/analysis/POLREF/RBNumber/RB2410236/"  
        files = [r"MET20240407-02-0.7T", r"MET20240407-02-0.2T", r"MET20240407-02-0.01T"]
        names = ["NB, I = 68 A", "NB, I = 20 A", "NB, I = 1 A"]

    elif code == 'TI_TA20_ML':
        path = r"/data/analysis/POLREF/RBNumber/RB2510555/"
        files = [r"TI_Ta20_ML_RT_68Amps", r"TI_Ta20_ML_RT_20Amps", r"TI_Ta20_ML_RT_10Amps", r"TI_Ta20_ML_RT_5Amps" , r"TI_Ta20_ML_RT_1Amp"]
        names = ["Ta(20 A), I = 68 A", "Ta(20 A), I = 20 A", "Ta(20 A), I = 10 A", "Ta(20 A), I = 5 A", "Ta(20 A), I = 1 A"]

    else:

        raise NameError(f'Sample Code: {code} does not exist')

        return()

    probes = []

    for f, n in zip(files, names):

        probe = load_probe_polref(filename= f, angle=theta, dQoQ=0.03, name= n,
                                    path = path, pol_mode="pnr", intensity=0.211,
                                    background=1e-7, back_reflectivity=False, sample_broadening=0.0)


        probe.pp.intensity.range(1e-1, 2)
        probe.pp.background.range(1e-9, 1e-5)
        probe.pp.sample_broadening.range(-(dQoQ*theta), 0.03)

        probes.append(probe)

    return(probes)

