import numpy as np
from scipy import interpolate

class QIE:
    def __init__(self, q, i, e=None, name="", color=None):
        """
        Initialize a QIE object.

        :param q: Array-like, the q values.
        :param i: Array-like, the i values.
        :param e: Array-like, the e values (optional).
        :param name: String, name of the QIE object (optional).
        :param color: Color for plotting (optional).
        """
        self.q = np.array(q)
        self.i = np.array(i)
        self.e = np.array(e) if e is not None else None
        self.name = name
        self.color = color

    def get_values(self):
        """
        Get the values of q, i, and e.

        :return: Tuple (q, i, e)
        """
        return self.q, self.i, self.e

    def get_qi(self):
        """
        Get the values of q and i.

        :return: Tuple (q, i)
        """
        return self.q, self.i

    def get_fill_between(self):
        """
        Get the values for fill_between plotting.

        :return: Tuple (q, i-e, i+e)
        """
        return self.q, self.i - self.e, self.i + self.e

    def set_q(self, q):
        """
        Set the q values.

        :param q: Array-like, new q values.
        """
        self.q = np.array(q)

    def set_i(self, i):
        """
        Set the i values.

        :param i: Array-like, new i values.
        """
        self.i = np.array(i)

    def clip_qmin(self, qmin):
        """
        Clip the values where q is greater than or equal to qmin.

        :param qmin: Minimum q value for clipping.
        :return: New QIE object with clipped values.
        """
        mask = self.q >= qmin
        return QIE(self.q[mask], self.i[mask], self.e[mask] if self.e is not None else None)

    def clip_qmax(self, qmax):
        """
        Clip the values where q is less than or equal to qmax.

        :param qmax: Maximum q value for clipping.
        :return: New QIE object with clipped values.
        """
        mask = self.q <= qmax
        return QIE(self.q[mask], self.i[mask], self.e[mask] if self.e is not None else None)

    def get_plot_kwargs(self):
        """
        Get keyword arguments for plotting.

        :return: Dictionary with plotting kwargs.
        """
        return {'label': self.name, 'color': self.color}

    def __repr__(self):
        """
        String representation of the QIE object.

        :return: String representation.
        """
        if self.name:
            return (f'{self.name} : nb values={len(self.q)}\n'
                    f'q={self.q}\n'
                    f'i={self.i}\n'
                    f'e={self.e}')
        else:
            return (f'nb values : {len(self.q)}\n'
                    f'q={self.q}\n'
                    f'i={self.i}\n'
                    f'e={self.e}')

def merge(qie0, qie1, cut1, cut2):
    """
    Merge two QIE curves by averaging between cut1 and cut2.

    :param qie0: First QIE object.
    :param qie1: Second QIE object.
    :param cut1: First cut value.
    :param cut2: Second cut value.
    :return: Merged QIE object.
    """
    if qie0.q[0] > qie1.q[0]:
        qie0, qie1 = qie1, qie0
    if cut1 > cut2:
        cut1, cut2 = cut2, cut1

    qie_final = qie0.clip_qmax(cut1)
    qie2 = qie0.clip_qmin(cut1).clip_qmax(cut2)
    qie3 = qie1.clip_qmin(cut1).clip_qmax(cut2)
    qie4 = qie1.clip_qmin(cut2)

    q2, i2, e2 = qie2.get_values()
    q3, i3, e3 = qie3.get_values()
    q_final, i_final, e_final = qie_final.get_values()

    newf = interpolate.interp1d(q2, i2, kind='linear', bounds_error=False, fill_value='extrapolate')
    newfe = interpolate.interp1d(q2, e2, kind='linear', bounds_error=False, fill_value='extrapolate')
    newi = newf(q3)
    newe = newfe(q3)

    q_final = np.concatenate((q_final, q3, qie4.q))
    i_final = np.concatenate((i_final, (newi + i3) / 2, qie4.i))
    e_final = np.concatenate((e_final, (newe + e3), qie4.e))

    return QIE(q_final, i_final, e_final)

def concatenate(qie0, qie1):
    """
    Concatenate two QIE curves.

    :param qie0: First QIE object.
    :param qie1: Second QIE object.
    :return: Concatenated QIE object.
    """
    q_final, i_final, e_final = qie0.get_values()
    q4, i4, e4 = qie1.get_values()

    q_final = np.concatenate((q_final, q4))
    i_final = np.concatenate((i_final, i4))
    e_final = np.concatenate((e_final, e4))

    return QIE(q_final, i_final, e_final)

def raw_subtract(qie0, qie1):
    """
    Subtract two QIE curves and add the errors.

    :param qie0: First QIE object.
    :param qie1: Second QIE object.
    :return: Resulting QIE object after subtraction.
    """
    q0, i0, e0 = qie0.get_values()
    q1, i1, e1 = qie1.get_values()
    return QIE(q0, i0 - i1, e0 + e1)

def smart_subtract(qie0,qie1):
    """
    considering q are different
    qie1 is solvent
    """
    q0, i0, e0 = qie0.get_values()
    qSolvent, iSolvent, eSolvent = qie1.get_values()

    newSolventFunction = interpolate.interp1d(qSolvent, iSolvent, kind='linear',
                                          bounds_error=0)  # interpolation for i
    new_iSolvent = newSolventFunction(q0)
    if eSolvent is not None:
        newErrFunction = interpolate.interp1d(qSolvent, eSolvent, kind='linear',
                                          bounds_error=0)  # interpolation for i
        new_eSolvent = newErrFunction(q0)
    new_qSolvent = q0

    iFinal = i0 - new_iSolvent
    if (eSolvent is not None) and (e0 is not None):
        errFinal = e0 + new_eSolvent
    else:
        errFinal=None
    return QIE(new_qSolvent,iFinal,errFinal)