import os
import csv
import numpy as np
from tqdm import tqdm
from matplotlib import pyplot as plt

from .Bind import Bind
from .utils import readBlockData, get_hist


class Calibration(object):
    """Calibrate the detector according to the calibration data"""

    def __init__(self):
        pass

    def __call__(self, file1, file2, bias_b, delta_e, n=6, m=8):
        """Calibration

        Parameters
        ----------
        file1/2 : str
            data file path of energy 1/2
        bias_b : float
            bias $b = b_1 + b_2$, in MeV
        delta_e : float
            delta energy between two beams, in MeV
        n : int, optional
            number of blocks, default 6
        m : int, optional
            number of binds, default 8
        """
        # Initialization
        self.n, self.m = n, m
        self.binds = [[Bind(i, j) for j in range(m)] for i in range(n)]

        # Read Data
        file_list = csv.reader(open(file1, "r"))
        pbar = tqdm(desc="Read Data E1", total=len(open(file1, "r").readlines()))
        for row in file_list:
            pn = int(row[1])
            ldata, rdata = readBlockData(row[0], int(row[5]), pn, m)
            for i in range(m):
                bind = self.binds[pn][i]
                bind.add_data(0, ldata[i], rdata[i], float(row[4]))
            pbar.update(1)
        pbar.close()

        file_list = csv.reader(open(file2, "r"))
        pbar = tqdm(desc="Read Data E2", total=len(open(file2, "r").readlines()))
        for row in file_list:
            pn = int(row[1])
            ldata, rdata = readBlockData(row[0], int(row[5]), pn, m)
            for i in range(m):
                bind = self.binds[pn][i]
                bind.add_data(1, ldata[i], rdata[i], float(row[4]))
            pbar.update(1)
        pbar.close()

        # Data preprocessing
        pbar = tqdm(desc="Bind Process", total=n * m)
        for i in range(n):
            for j in range(m):
                bind: Bind = self.binds[i][j]
                bind.slash()
                bind.get_line()
                bind.get_kb(bias_b, delta_e)
                pbar.update(1)
        pbar.close()

        # Fit
        pbar = tqdm(desc="Bind Fit", total=n * m)
        for i in range(n):
            for j in range(m):
                bind: Bind = self.binds[i][j]
                bind.get_peak_center()
                bind.fit_px()
                pbar.update(1)
        pbar.close()

    def draw_fit(self, path):
        """Draw fit result

        Parameters
        ----------
        path : str
            save folder, there must be `FIT-LINE`, `GMM`, and `PEAK` 3 subfolders.
        """
        pbar = tqdm(desc="Draw Fit Figure", total=self.n * self.m)
        for i in range(self.n):
            for j in range(self.m):
                bind: Bind = self.binds[i][j]
                bind.draw_fit_line(os.path.join(path, "FIT-LINE", bind.name + ".png"))
                bind.draw_cluster(os.path.join(path, "GMM", bind.name + ".png"))
                bind.draw_peak(os.path.join(path, "PEAK", bind.name + ".png"))

                pbar.update(1)
        pbar.close()

    def draw_check(self, path="Check.png"):
        """Draw check figure

        Parameters
        ----------
        path : str, optional
            save path
        """
        pbar = tqdm(desc="Draw Check Figure", total=self.n * self.m)
        fig = plt.figure(figsize=(24, 24), dpi=200)
        ax1 = fig.add_subplot(3, 1, 1)
        ax2 = fig.add_subplot(3, 1, 2)
        ax3 = fig.add_subplot(3, 1, 3)
        peaks = np.array([])

        for i in range(self.n):
            engs = np.array([])
            for j in range(self.m):
                bind = self.binds[i][j]
                peaks = np.hstack((np.unique(bind.px[0]), peaks))

                eng = bind.predict_energy(bind.x[0], bind.y[0])
                pX = bind.predict_px(bind.x[0], bind.y[0])
                count, center = get_hist(pX, delta=0.5)
                engs = np.hstack((engs, eng))

                ax2.scatter(pX, eng, s=0.1, color="k")
                ax3.scatter(center, count + 3000 * (7 - j), s=0.5, color="k")

                pbar.update(1)
            count, center = get_hist(engs, delta=0.01)
            ax1.plot(center, count)

        peaks = np.unique(peaks)
        for x in peaks:
            ax3.vlines(x, 0, 3000 * self.m, color="gray", linestyles="dashed")
        for j in range(self.m):
            ax3.hlines(
                3000 * j,
                (np.min(peaks) // 50) * 50,
                (np.max(peaks) // 50 + 1) * 50,
                color="r",
                linestyles="dashdot",
            )

        ax1.set_xlabel("Energy (MeV)")
        ax1.set_ylabel("Count per bin")
        ax2.set_xlabel("x (mm)")
        ax2.set_ylabel("Energy (MeV)")
        ax3.set_xlabel("x (mm)")
        ax3.set_ylabel("Count per bin")

        fig.savefig(path, facecolor="w", transparent=False)
        plt.close()
        pbar.close()

    def save(self, path="coef.csv"):
        """Save coefficient to file

        Parameters
        ----------
        path : str, optional
            save path
        """
        f = open(path, "w")
        for i in range(self.n):
            for j in range(self.m):
                bind = self.binds[i][j]
                f.writelines(
                    "{:d},{:d},{:.9f},{:.9f},{:.9f},{:.9f},{:.9f}\n".format(
                        i, j, bind.k1, bind.k2, bind.b, bind.L, bind.C
                    )
                )