import numpy as np
import matplotlib.pyplot as plt
from astropy.modeling import models, fitting
from sklearn.mixture import GaussianMixture
from sklearn.linear_model import LinearRegression
from .utils import get_hist

class Bind(object):
    flag = 0
    L, C = 40, 0
    k1, k2, b = 0, 0, 0
    K1, C1, K2, C2 = 0, 0, 0, 0

    def __init__(self, n, m):
        self.n, self.m = n, m

    def add_data(self, E, data):
        if self.flag == 0:
            self.flag = 1
            self.E1 = E
            self.x1, self.y1 = data[:, 0].reshape(-1, 1), data[:, 1].reshape(-1, 1)
        else:
            self.E2 = E
            self.x2, self.y2 = data[:, 0].reshape(-1, 1), data[:, 1].reshape(-1, 1)

    def fit1(self):
        self.reg1 = LinearRegression()
        self.reg1.fit(self.x1, self.y1)
        self.K1 = self.reg1.coef_[0][0]
        self.C1 = self.reg1.intercept_[0]

        self.reg2 = LinearRegression()
        self.reg2.fit(self.x2, self.y2)
        self.K2 = self.reg2.coef_[0][0]
        self.C2 = self.reg2.intercept_[0]

    def filter(self):
        self.fit1()
        ny1 = self.y1 - self.reg1.predict(self.x1)
        ny2 = self.y2 - self.reg2.predict(self.x2)

        idx1 = np.where(ny1 <= ny1.std())[0]
        idx2 = np.where(ny2 <= ny2.std())[0]

        self.x1 = self.x1[idx1]
        self.y1 = self.y1[idx1]
        self.x2 = self.x2[idx2]
        self.y2 = self.y2[idx2]

    def solve(self):
        self.k1 = (self.E1 - self.E2) / (self.C1 - self.C2)
        self.k2 = -(self.K1 + self.K2) * self.k1 / 2
        self.b = (self.E1 - self.C1 * self.k1 + self.E2 - self.C2 * self.k1) / 2

    def split(self, n=7):
        self.clusters = []
        data = np.concatenate((self.x1, self.y1), axis=1)

        model = GaussianMixture(n_components=n)
        model.fit(data)

        ny = model.predict(data)
        for i in np.unique(ny):
            idx = np.where(ny == i)[0]
            self.clusters.append(data[idx])

        self.clusters = np.array(self.clusters, dtype=object)
        self.clusters = sorted(self.clusters, key=lambda x: len(x))

    def fit2(self):
        self.fx = []
        self.fy = []
        fitter = fitting.LevMarLSQFitter()

        for data in self.clusters:
            x1, x2 = data[:, 0], data[:, 1]
            c1, e1 = get_hist(x1)
            c2, e2 = get_hist(x2)
            modelx1 = models.Gaussian1D(amplitude=c1.max(), mean=x1.mean(), stddev=x1.std())
            modelx2 = models.Gaussian1D(amplitude=c2.max(), mean=x2.mean(), stddev=x2.std())
            fitted_model1 = fitter(modelx1, e1, c1)
            fitted_model2 = fitter(modelx2, e2, c2)
            self.fx.append(fitted_model1.mean.value)
            self.fy.append(fitted_model2.mean.value)

        idx = np.argsort(self.fx)

    def fit3(self, y):
        self.fx = np.sort(self.fx)[::-1].reshape(-1, 1)
        self.reg3 = LinearRegression()
        self.reg3.fit(self.fx, y)
        self.K3 = self.reg3.coef_[0]
        self.C3 = self.reg3.intercept_

    def draw_fit1(self, title):
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(1, 1, 1)
        ax.scatter(self.x1, self.y1, s=0.1, c='black', label=r'$E_1$')
        ax.scatter(self.x2, self.y2, s=0.1, c='dimgray', label=r'$E_2$')
        ax.plot(self.x1, self.reg1.predict(self.x1), c='red', label=r'$x_2={:.4f}x_1+{:.4f},\ R^2={:.5f}$'.format(self.K1, self.C1, self.RSquare1))
        ax.plot(self.x2, self.reg2.predict(self.x2), c='orangered', label=r'$x_2={:.4f}x_1+{:.4f},\ R^2={:.5f}$'.format(self.K2, self.C2, self.RSquare2))
        ax.legend()
        fig.savefig(title, facecolor='w', transparent=False)
        plt.close()

    def draw_split(self, title):
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(1, 1, 1)
        for data in self.clusters:
            ax.scatter(data[:, 0], data[:, 1], s=0.1)
        fig.savefig(title, facecolor='w', transparent=False)
        plt.close()

    def draw_fit2(self, title):
        k, n = 1, len(self.clusters)
        fig = plt.figure(figsize=(8, 10))
        for data in self.clusters:
            ax = fig.add_subplot(n, 1, k)
            x1, x2 = data[:, 0], data[:, 1]
            c1, e1 = get_hist(x1)
            c2, e2 = get_hist(x2)
            ax.scatter(e1, c1, s=0.1)
            ax.scatter(e2, c2, s=0.1)
            k += 1
        plt.tight_layout()
        fig.savefig(title, facecolor='w', transparent=False)
        plt.close()

    @property
    def name(self):
        return '{:.1f}-{:.1f}-{:d}-{:d}'.format(self.E1, self.E2, self.n, self.m)

    @property
    def RSquare1(self):
        return self.reg1.score(self.x1, self.y1)

    @property
    def RSquare2(self):
        return self.reg2.score(self.x2, self.y2)