from scipy.optimize import fsolve
from autograd import grad
from .binding_system import BindingSystem
from .lagrange_binding_system_factory import LagrangeBindingSystemFactory
import numpy as np

# 1:1 binding - see https://stevenshave.github.io/pybindingcurve/simulate_1to1.html
def system01_lagrange(p, l, kdpl):
    def F(X):  # Augmented Lagrange function
        response = (X[0] * X[1]) / kdpl
        return (
            response - X[2] * (p - (response + X[0])) - X[3] * (l - (response + X[1]))
        )

    dfdL = grad(F)  # Gradients of the Lagrange function
    pf, lf, lam1, lam2 = fsolve(dfdL, [p, l, 1.0, 1.0])
    return {"pf": pf, "lf": lf, "pl": (pf * lf) / kdpl}


# 1:1:1 competition - see https://stevenshave.github.io/pybindingcurve/simulate_competition.html
def system02_lagrange(p, l, i, kdpl, kdpi):
    def F(X):  # Augmented Lagrange function
        response = (X[0] * X[1]) / kdpl
        constraint1 = p - (response + X[0] + (X[0] * X[2] / kdpi))
        constraint2 = l - (response + X[1])
        constraint3 = i - (X[2] + (X[0] * X[2] / kdpi))
        return response - X[3] * constraint1 - X[4] * constraint2 - X[5] * constraint3

    dfdL = grad(F, 0)  # Gradients of the Lagrange function
    pf, lf, inhf, lam1, lam2, lam3 = fsolve(dfdL, [p, l, i, 1.0, 1.0, 1.0])
    return {"pf": pf, "lf": lf, "pl": (pf * lf) / kdpl}


# Homodimer formation - see https://stevenshave.github.io/pybindingcurve/simulate_homodimerformation.html
def system03_lagrange(p, kdpp):
    def F(X):  # Augmented Lagrange function
        response = (X[0] * X[0]) / kdpp
        constraint1 = p - (response * 2 + X[0])
        return response - X[1] * constraint1

    dfdL = grad(F, 0)  # Gradients of the Lagrange function
    pf, lam1 = fsolve(dfdL, [p, 1.0])
    return {"pf": pf, "pp": ((pf * pf) / kdpp)}


# Homodimer breaking - see https://stevenshave.github.io/pybindingcurve/simulate_homodimerbreaking.html
def system04_lagrange(p, i, kdpp, kdpi):
    def F(X):  # Augmented Lagrange function
        pp = (X[0] * X[0]) / kdpp
        pi = (X[0] * X[1]) / kdpi
        constraint1 = p - (pp * 2 + pi + X[0])
        constraint2 = i - (pi + X[1])
        return pp - X[2] * constraint1 - X[3] * constraint2

    dfdL = grad(F, 0)  # Gradients of the Lagrange function
    pf, i_f, lam1, lam2 = fsolve(dfdL, [p, i, 1.0, 1.0])
    return {"pf": pf, "if": i_f, "pp": ((pf * pf) / kdpp), "pi": (pf * i_f) / kdpi}


# 1:2 binding
def system12_lagrange(p, l, kdpl1, kdpl2):
    def F(X):  # Augmented Lagrange function
        pf = X[0]
        lf = X[1]
        pl1 = pf * lf / kdpl1
        pl2 = pf * lf / kdpl2
        pl12 = (pl1 * lf + pl2 * lf) / (kdpl1 + kdpl2)
        constraint1 = p - (pf + pl1 + pl2 + pl12)
        constraint2 = l - (lf + pl1 + pl2 + 2 * pl12)
        return pl12 - X[2] * constraint1 - X[3] * constraint2

    dfdL = grad(F, 0)  # Gradients of the Lagrange function

    pf, lf, lam1, lam2 = fsolve(dfdL, [p, l] + [1.0] * 2)
    pl1 = pf * lf / kdpl1
    pl2 = pf * lf / kdpl2
    pl12 = (pl1 * lf + pl2 * lf) / (kdpl1 + kdpl2)
    return {"pf": pf, "lf": lf, "pl1": pl1, "pl2": pl2, "pl12": pl12}


# 1:3 binding
def system13_lagrange(p, l, kdpl1, kdpl2, kdpl3):
    def F(X):  # Augmented Lagrange function
        pf = X[0]
        lf = X[1]
        pl1 = pf * lf / kdpl1
        pl2 = pf * lf / kdpl2
        pl3 = pf * lf / kdpl3
        pl12 = (pl1 * lf + pl2 * lf) / (kdpl1 + kdpl2)
        pl13 = (pl1 * lf + pl3 * lf) / (kdpl1 + kdpl3)
        pl23 = (pl2 * lf + pl3 * lf) / (kdpl2 + kdpl3)
        pl123 = (pl12 * lf + pl13 * lf + pl23 * lf) / (kdpl1 + kdpl2 + kdpl3)
        constraint1 = p - (pf + pl1 + pl2 + pl3 + pl12 + pl13 + pl23 + pl123)
        constraint2 = l - (
            lf + pl1 + pl2 + pl3 + 2 * (pl12 + pl13 + pl23) + 3 * (pl123)
        )
        nonzero_constraint = (constraint1 - abs(constraint1)) - (
            constraint2 - abs(constraint2)
        )
        return (
            pl123 - X[2] * constraint1 - X[3] * constraint2 - X[4] * nonzero_constraint
        )

    dfdL = grad(F, 0)  # Gradients of the Lagrange function
    pf, lf, lam1, lam2, lam3 = fsolve(dfdL, [p, l] + [1.0] * 3)
    pl1 = pf * lf / kdpl1
    pl2 = pf * lf / kdpl2
    pl3 = pf * lf / kdpl3
    pl12 = (pl1 * lf + pl2 * lf) / (kdpl1 + kdpl2)
    pl13 = (pl1 * lf + pl3 * lf) / (kdpl1 + kdpl3)
    pl23 = (pl2 * lf + pl3 * lf) / (kdpl2 + kdpl3)
    pl123 = (pl12 * lf + pl13 * lf + pl23 * lf) / (kdpl1 + kdpl2 + kdpl3)
    return {
        "pf": pf,
        "lf": lf,
        "pl1": pl1,
        "pl2": pl2,
        "pl3": pl3,
        "pl12": pl12,
        "pl13": pl13,
        "pl23": pl23,
        "pl123": pl123,
    }


class System_lagrange_one_to_one__pl(BindingSystem):
    """
    Lagrange 1:1 binding system

    Class defines 1:1 binding, readout is PL
    See https://stevenshave.github.io/pybindingcurve/simulate_1to1.html
    """

    def __init__(self):
        super().__init__(system01_lagrange)
        self.default_readout = "pl"

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / parameters["l"]
                )
        else:
            return super().query(parameters)


class System_lagrange_competition_pl(BindingSystem):
    """
    Lagrange 1:1:1 competition binding system

    Class defines 1:1:1 competition, readout is PL
    See https://stevenshave.github.io/pybindingcurve/simulate_competition.html
    """

    def __init__(self):
        super().__init__(system02_lagrange)
        self.default_readout = "pl"

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / parameters["l"]
                )
        else:
            return super().query(parameters)


class System_lagrange_homodimerformation__pp(BindingSystem):
    """
    Lagrange homodimer formation system

    Class defines homodimer formation, readout is PP
    See https://stevenshave.github.io/pybindingcurve/simulate_homodimerformation.html
    """

    def __init__(self):
        super().__init__(system03_lagrange, False)
        self.default_readout = "pp"

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / (parameters["p"] / 2.0)
                )
        else:
            return super().query(parameters)


class System_lagrange_homodimerbreaking__pp(BindingSystem):
    """
    Lagrange homodimer breaking system

    Class defines homodimer breaking, readout is PP
    See https://stevenshave.github.io/pybindingcurve/simulate_homodimerbreaking.html
    """

    def __init__(self):
        super().__init__(system04_lagrange, False)
        self.default_readout = "pp"

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / (parameters["p"] / 2.0)
                )
        else:
            return super().query(parameters)


class System_lagrange_homodimerbreaking__pl(BindingSystem):
    """
    Lagrange homodimer breaking system

    Class defines homodimer breaking, readout is PL
    See https://stevenshave.github.io/pybindingcurve/simulate_homodimerbreaking.html
    """

    def __init__(self):
        super().__init__(system04_lagrange, False)
        self.default_readout = "pl"

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / (parameters["p"] / 2.0)
                )
        else:
            return super().query(parameters)


class System_lagrange_1_to_2__pl12(BindingSystem):
    """
    Lagrange 1:2 binding system.

    Class defines 1:2 protein:ligand binding, readout is PL12, meaning protein
    with 2 ligands
    See https://stevenshave.github.io/pybindingcurve/simulate_homodimerbreaking.html
    """

    def __init__(self):
        super().__init__(system12_lagrange, False)
        self.default_readout = "pl12"

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / parameters["l"]
                )
        else:
            return super().query(parameters)


class System_lagrange_1_to_3__pl123(BindingSystem):
    """
    Lagrange 1:2 binding system.

    Class defines 1:2 protein:ligand binding, readout is PL123, meaning protein
    with 3 ligands
    See https://stevenshave.github.io/pybindingcurve/simulate_homodimerbreaking.html
    """

    def __init__(self):
        super().__init__(system13_lagrange, False)
        self.default_readout = "pl123"

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / parameters["l"]
                )
        else:
            return super().query(parameters)


class System_lagrange_custom(BindingSystem):
    """
    Lagrange custom binding system

    Class uses LagrangeBindingSystemFactory to make a custom lagrange function.
    """

    def __init__(self, system_string):
        custom_system = LagrangeBindingSystemFactory(system_string)
        super().__init__(custom_system.binding_function)
        self.default_readout = custom_system.default_readout

    def query(self, parameters: dict):
        if self._are_ymin_ymax_present(parameters):
            parameters_no_min_max = self._remove_ymin_ymax_keys_from_dict_return_new(
                parameters
            )
            value = super().query(parameters_no_min_max)
            with np.errstate(divide="ignore", invalid="ignore"):
                return np.nan_to_num(
                    parameters["ymin"]
                    + ((parameters["ymax"] - parameters["ymin"]) * value)
                    / parameters["l"]
                )
        else:
            return super().query(parameters)
