# Define the path to the data file relative to the current script location
import os

import numpy as np
from scipy.interpolate import CubicSpline

from cosmo_wap.lib import utils
from cosmo_wap.lib.luminosity_funcs import (
    BGSLuminosityFunction,
    LBGLuminosityFunction,
    Model1LuminosityFunction,
    Model3LuminosityFunction,
)


class SurveyParams:
    def get(self, cosmo, survey):
        """
        Initialize and get survey parameters for some set surveys
        """
        return getattr(self, survey)(cosmo)

    # ok want to inherit this function to update variables - could use dataclasses
    class SurveyBase:
        def update(self, **kwargs):
            """update survey class parameters-
            Input dictionary with parameters you want to change."""
            new_self = utils.copy(self)

            for key, value in kwargs.items():
                if hasattr(new_self, key):
                    setattr(new_self, key, value)
            return new_self

        def compute_luminosity(self, LF, cut, zz, need_hod=False):
            """Get biases from given luminosity function and magnitude/luminosity cut

            LF: Luminosity function class
            cut: Magnitude/flux cut
            zz: redshift
            need_hod: if True, defer bias computation to HOD later
            Returns:
            Object with Q,be and n_g defined"""
            self.cut = cut
            self.need_hod = need_hod
            if not need_hod:
                self.Q = CubicSpline(zz, LF.get_Q(cut, zz))
                self.be = CubicSpline(zz, LF.get_be(cut, zz))
                self.n_g = CubicSpline(zz, LF.number_density(cut, zz))
                if hasattr(LF, "get_b_1"):  # then also get linear bias from fits in Table. 2 1909.12069
                    self.b_1 = CubicSpline(zz, LF.get_b_1(cut, zz))
            return self

        def BF_split(self, split):
            """
            Split is magnitude or flux cut at which we seperate our samples
            Call from survey class with a luminosity function.
            """
            if not hasattr(self, "LF"):
                raise ValueError("No luminosity function - use survey with defined luminosity function!")

            self.split = split  # store for later
            self.bright = utils.copy(self)  # bright
            self.faint = utils.copy(self)  # faint
            if not self.need_hod:
                """
                We already have total sample biases.
                n_F = n_T - n_B
                bF = n_T b_T − n_B b_B /nF
                Q_F = n_T/(n_T - n_B) Q_T - n_B/(n_T - n_B)*Q_B
                be_F = d ln(n_T - n_B)/ d ln (1 + z)
                """
                zz = self.zz
                # define total survey biases
                Q_T = self.Q(zz)
                n_T = self.n_g(zz)

                Q_B = self.LF.get_Q(split, zz)
                be_B = self.LF.get_be(split, zz)
                n_B = self.LF.number_density(split, zz)

                self.bright.Q = CubicSpline(zz, Q_B)
                self.bright.be = CubicSpline(zz, be_B)
                self.bright.n_g = CubicSpline(zz, n_B)

                # so for faint
                self.faint.Q = CubicSpline(zz, n_T / (n_T - n_B) * Q_T - n_B / (n_T - n_B) * Q_B)
                self.faint.be = CubicSpline(zz, np.gradient(np.log(n_T - n_B), np.log(1 + zz)))
                self.faint.n_g = CubicSpline(zz, n_T - n_B)

                # then for linear bias if we can use semi-analytical fit from 1909.12069
                if hasattr(self.LF, "get_b_1"):
                    b_T = self.b_1(zz)
                    b_B = self.LF.get_b_1(split, zz)
                    self.bright.b_1 = CubicSpline(zz, b_B)
                    self.faint.b_1 = CubicSpline(zz, (n_T * b_T - n_B * b_B) / (n_T - n_B))

                return [self.bright, self.faint]  # two tracers defined
            else:
                return self

        def load_SKAO_data(self):
            module_dir = os.path.dirname(os.path.abspath(__file__))
            self.SKAO1Data = np.loadtxt(os.path.join(module_dir, "data_library/SKAO1Data.txt"))
            self.SKAO2Data = np.loadtxt(os.path.join(module_dir, "data_library/SKAO2Data.txt"))

    class Euclid(SurveyBase):
        def __init__(self, cosmo, fitting=False, model3=True, F_c=None):
            self.cosmo = cosmo
            self.b_1 = lambda xx: 0.9 + 0.4 * xx
            self.f_sky = 15000 / 41253
            self.z_range = [0.9, 1.8]  # get zmin and zmax

            if fitting:
                self.be = lambda xx: -7.29 + 0.470 * xx + 1.17 * xx**2 - 0.290 * xx**3  # euclid_data[:,2]
                self.Q = lambda xx: 0.583 + 2.02 * xx - 0.568 * xx**2 + 0.0411 * xx**3
                self.n_g = lambda zz: 0.0193 * zz ** (-0.0282) * np.exp(-2.81 * zz)
            else:
                self.zz = np.linspace(self.z_range[0], self.z_range[1], 100)
                # from lumnosity function
                if model3:
                    if F_c is None:  # set defualt values - this one agrees with fitting functions above
                        F_c = 2e-16
                    self.LF = Model3LuminosityFunction(cosmo)
                else:
                    if F_c is None:
                        F_c = 3e-16
                    self.LF = Model1LuminosityFunction(cosmo)

                self.compute_luminosity(self.LF, F_c, self.zz)

    class Roman(SurveyBase):
        def __init__(self, cosmo, model3=False, F_c=None):
            self.cosmo = cosmo
            self.b_1 = lambda xx: 0.9 + 0.4 * xx
            self.f_sky = 2000 / 41253
            self.z_range = [0.5, 2.0]  # get zmin and zmax
            self.zz = np.linspace(self.z_range[0], self.z_range[1], 100)

            # from lumnosity function
            if model3:
                if F_c is None:  # set defualt values
                    F_c = 1e-16
                self.LF = Model3LuminosityFunction(cosmo)
            else:
                if F_c is None:
                    F_c = 1e-16
                self.LF = Model1LuminosityFunction(cosmo)

            self.compute_luminosity(self.LF, F_c, self.zz)

    class BGS(SurveyBase):
        def __init__(self, cosmo, m_c=20, flag="HOD"):
            self.cosmo = cosmo
            self.b_1 = lambda xx: 1.34 / cosmo.scale_independent_growth_factor(xx)
            self.z_range = [0.05, 0.6]
            self.f_sky = 15000 / 41253

            # default
            self.be = lambda xx: -2.25 - 4.02 * xx + 0.318 * xx**2 - 14.6 * xx**3
            self.Q = lambda xx: 0.282 + 2.36 * xx + 2.27 * xx**2 + 11.1 * xx**3
            self.n_g = lambda zz: 0.023 * zz ** (-0.471) * np.exp(-5.17 * zz) - 0.002  # fitting from Maartens
            if flag == "LF":
                # from lumnosity function
                self.zz = np.linspace(self.z_range[0], self.z_range[1], 100)
                self.LF = BGSLuminosityFunction(cosmo)
                self.compute_luminosity(self.LF, m_c, self.zz)
            elif flag == "HOD":
                self.zz = np.linspace(self.z_range[0], self.z_range[1], 100)
                self.LF = BGSLuminosityFunction(cosmo)
                self.compute_luminosity(self.LF, m_c, self.zz, need_hod=True)  # compute biases with HOD later

    class MegaMapper(SurveyBase):
        def __init__(self, cosmo, m_c=24.5):
            self.cosmo = cosmo
            self.A = -0.98 * (m_c - 25) + 0.11  # from Eq.(2.7) 1904.13378v2
            self.B = 0.12 * (m_c - 25) + 0.17
            self.b_1 = lambda xx: (
                self.A * (1 + xx) + self.B * (1 + xx) ** 2
            )  # so linear bias for given apparent magnitude limit
            self.z_range = [2.1, 5]  # get zmin and zmax
            self.f_sky = 20000 / 41253

            # from lumnosity function
            self.LF = LBGLuminosityFunction(cosmo)
            self.zz = self.LF.z_values
            self.compute_luminosity(self.LF, m_c, self.zz)

    class SKAO1(SurveyBase):
        def __init__(self, cosmo):
            self.cosmo = cosmo
            self.load_SKAO_data()
            self.b_1 = lambda xx: 0.616 * np.exp(1.017 * xx)
            self.z_range = [self.SKAO1Data[:, 0][0], self.SKAO1Data[:, 0][-1]]
            self.be = CubicSpline(self.SKAO1Data[:, 0], self.SKAO1Data[:, 4])
            self.Q = CubicSpline(self.SKAO1Data[:, 0], self.SKAO1Data[:, 3])
            self.n_g = CubicSpline(self.SKAO1Data[:, 0], self.SKAO1Data[:, 2])  # fitting from Maartens
            self.f_sky = 5000 / 41253

    class SKAO2(SurveyBase):
        def __init__(self, cosmo):
            self.cosmo = cosmo
            self.load_SKAO_data()
            self.b_1 = lambda xx: 0.554 * np.exp(0.783 * xx)
            self.z_range = [self.SKAO2Data[:, 0][0], self.SKAO2Data[:, 0][-1]]
            self.be = CubicSpline(self.SKAO2Data[:, 0], self.SKAO2Data[:, 4])
            self.Q = CubicSpline(self.SKAO2Data[:, 0], self.SKAO2Data[:, 3])
            self.n_g = CubicSpline(self.SKAO2Data[:, 0], self.SKAO2Data[:, 2])  # fitting from Maartens
            self.f_sky = 30000 / 41253

    class DM_part(SurveyBase):
        def __init__(self, cosmo):
            self.cosmo = cosmo
            self.b_1 = lambda xx: 1 + 0 * xx  # for dark matter particles
            self.b_2 = lambda xx: 0 * xx
            self.g_2 = lambda xx: 0 * xx
            self.z_range = [0.01, 5]
            self.be = lambda xx: 0 * xx
            self.Q = lambda xx: 0 * xx
            self.n_g = lambda xx: 1e5 + 0 * xx
            self.f_sky = 1

    class InitNew(SurveyBase):  # for adding new surveys... # so could add some sort of dict unpacking method?
        def __init__(self, cosmo):
            self.cosmo = cosmo
            self.b_1 = lambda xx: 1 + 0 * xx  # for dark matter particles
            self.z_range = [0.01, 5]
            self.be = lambda xx: 0 * xx
            self.Q = lambda xx: 0 * xx
            self.n_g = lambda xx: 1e5 + 0 * xx
            self.f_sky = 1


################################################################################


class SetSurveyFunctions:
    """Reads in survey specific params and calculates higher order ones if unprovided throguh set relations - if empty use default values.
    This is called once for each tracer
    """

    def __init__(self, survey_params, compute_bias=False):
        self.be = getattr(survey_params, "be", lambda xx: 0 * xx)
        self.Q = getattr(survey_params, "Q", lambda xx: 0 * xx)

        if not compute_bias:  # then just use default or assigned
            self.n_g = getattr(survey_params, "n_g", lambda xx: 1e5 + 0 * xx)  # no shot noise
            self.b_1 = getattr(survey_params, "b_1", lambda xx: np.sqrt(1 + xx))  # stupid linear bias model
            self.g_2 = getattr(
                survey_params, "g_2", lambda xx: -(2 / 7) * (self.b_1(xx) - 1)
            )  # from local langragian expression
            self.b_2 = getattr(
                survey_params,
                "b_2",
                lambda xx: (
                    0.412
                    - 2.143 * self.b_1(xx)
                    + 0.929 * self.b_1(xx) ** 2
                    + 0.008 * self.b_1(xx) ** 3
                    + 4 / 3 * self.g_2(xx)
                ),
            )  # quadratic fit: 1511.01096

            # Local PNG from universality
            class Loc:
                def __init__(self, parent):
                    delta_c = 1.686
                    bL10 = lambda xx: parent.b_1(xx) - 1
                    bL20 = lambda xx: parent.b_2(xx) - (8 / 21) * bL10(xx)

                    bL01 = lambda xx: 2 * delta_c * bL10(xx)
                    bL11 = lambda xx: 2 * (delta_c * bL20(xx) - bL10(xx))
                    bL02 = lambda xx: 4 * delta_c * (delta_c * bL20(xx) - 2 * bL10(xx))

                    self.b_01 = bL01
                    self.b_11 = lambda xx: bL01(xx) + bL11(xx)
                    self.b_02 = bL02

            self.loc = Loc(self)

        self.f_sky = getattr(survey_params, "f_sky", 1)
        self.z_range = getattr(survey_params, "z_range", np.linspace(0, 5, int(1e5)))

        self.betas = None
        self.deriv = {}

    def reset_cache(self):
        """Clear cached betas and derivatives — call when bias functions change."""
        self.betas = None
        self.deriv = {}
