import glob, sys
import uproot
import copy
from texttable import Texttable

import numpy as np

np.set_printoptions(threshold=sys.maxsize)

from ctapipe_io_magic import MAGICEventSource

from astropy import units as u

from ctapipe.instrument import CameraGeometry
from ctapipe.image import hillas_parameters

import pylab as plt

import argparse
import csv

from magicctapipe.utils import MAGIC_Badpixels

# from magicctapipe.utils import bad_pixel_treatment
from magicctapipe.utils import MAGIC_Cleaning


class bcolors:
    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


def file_reader(filename, max_row=1e30):
    datafile = open(filename, "r")
    datareader = csv.reader(datafile)
    data = []
    for ix, row in enumerate(datareader):
        if ix >= max_row:
            break
        if len(row) == 0:
            data.append(row)
        else:
            data.append(list(map(int, row[0][:-1].split(" "))))

    datafile.close()
    return data


def main():
    usage = "usage: %(prog)s [options] "
    description = "Run gtselect and gtmktime on one or more FT1 files.  "
    "Note that gtmktime will be skipped if no FT2 file is provided."
    parser = argparse.ArgumentParser(usage=usage, description=description)
    parser.add_argument("-m", "--mc", dest="mc", action="store_true")

    parser.add_argument("--base", dest="base", action="store_true")
    parser.add_argument("--nopix", dest="nopix", action="store_true")

    parser.add_argument("--test", dest="test", action="store_true")
    # parser.add_argument('--testnopix', dest='testnopix', action='store_true')

    args = parser.parse_args()

    tel_id = 1
    max_events = 1e40

    mars_camera = CameraGeometry.from_name("MAGICCamMars")

    neighbors = mars_camera.neighbor_matrix_sparse
    outermost = []
    for pix in range(mars_camera.n_pixels):
        if neighbors[pix].getnnz() < 5:
            outermost.append(pix)

    cleaning_config = dict(
        picture_thresh=6,
        boundary_thresh=3.5,
        max_time_off=4.5 * 1.64,
        max_time_diff=1.5 * 1.64,
        usetime=True,
        usesum=True,
        findhotpixels=True,
        test=False,
    )

    if args.base:
        base_dir = "/home/iwsatlas1/damgreen/CTA/MAGIC_Cleaning/star/Star_Base/"

        star_file = "%s/20191127_M1_05086952.001_I_CrabNebula-W0.40+035.root" % base_dir
        event_images_stardump = np.genfromtxt(
            "%s/event_images.txt" % base_dir, dtype=float, max_rows=max_events
        )
        event_times_stardump = np.genfromtxt(
            "%s/event_times.txt" % base_dir, dtype=float, max_rows=max_events
        )
        event_unsuitables_stardump = np.genfromtxt(
            "%s/event_unsuitable.txt" % base_dir, dtype=bool, max_rows=max_events
        )
        event_unmappeds_stardump = np.genfromtxt(
            "%s/event_unmapped.txt" % base_dir, dtype=bool, max_rows=max_events
        )

        event_cleanstep1 = file_reader(
            "%s/event_cleanstep1.txt" % base_dir, max_row=max_events
        )
        event_cleanstep2 = file_reader(
            "%s/event_cleanstep2.txt" % base_dir, max_row=max_events
        )
        event_cleanstep3 = file_reader(
            "%s/event_cleanstep3.txt" % base_dir, max_row=max_events
        )

        cleaning_config["findhotpixels"] = True

    if args.nopix:
        base_dir = "/home/iwsatlas1/damgreen/CTA/MAGIC_Cleaning/star/Star_No_HotPix/"

        star_file = "%s/20191127_M1_05086952.001_I_CrabNebula-W0.40+035.root" % base_dir
        event_images_stardump = np.genfromtxt(
            "%s/event_images.txt" % base_dir, dtype=float, max_rows=max_events
        )
        event_times_stardump = np.genfromtxt(
            "%s/event_times.txt" % base_dir, dtype=float, max_rows=max_events
        )
        event_unsuitables_stardump = np.genfromtxt(
            "%s/event_unsuitable.txt" % base_dir, dtype=bool, max_rows=max_events
        )
        event_unmappeds_stardump = np.genfromtxt(
            "%s/event_unmapped.txt" % base_dir, dtype=bool, max_rows=max_events
        )

        event_cleanstep1 = file_reader(
            "%s/event_cleanstep1.txt" % base_dir, max_row=max_events
        )
        event_cleanstep2 = file_reader(
            "%s/event_cleanstep2.txt" % base_dir, max_row=max_events
        )
        event_cleanstep3 = file_reader(
            "%s/event_cleanstep3.txt" % base_dir, max_row=max_events
        )

        cleaning_config["findhotpixels"] = False

    bad_pixels_config = dict(
        pedestalLevel=400, pedestalLevelVariance=4.5, pedestalType="FromExtractorRndm"
    )

    # Grab the processed star information
    hillas_array_list = [
        "MHillas.fSize",
        "MHillas.fWidth",
        "MHillas.fLength",
        "MRawEvtHeader.fStereoEvtNumber",
        "MNewImagePar.fNumCorePixels",
        "MNewImagePar.fNumUsedPixels",
    ]

    input_file = uproot.open(star_file)
    events = input_file["Events"].arrays(hillas_array_list)

    evtnum_arr = events[b"MRawEvtHeader.fStereoEvtNumber"]
    size_arr = events[b"MHillas.fSize"]
    width_arr = events[b"MHillas.fWidth"]
    length_arr = events[b"MHillas.fLength"]
    core_arr = events[b"MNewImagePar.fNumCorePixels"]
    used_arr = events[b"MNewImagePar.fNumUsedPixels"]

    star_size_arr = []
    ctapipe_size_arr = []

    star_length_arr = []
    ctapipe_length_arr = []

    star_width_arr = []
    ctapipe_width_arr = []

    magic_clean = MAGIC_Cleaning.magic_clean(mars_camera, cleaning_config)
    badpixel_calculator = MAGIC_Badpixels.MAGICBadPixelsCalc(config=bad_pixels_config)

    size_cut = 0
    AberrationCorrection = 1.0713

    if args.mc:
        """MC FILES"""
        mars_file_mask = "/remote/ceph/group/magic/MAGIC-LST/MCs/MAGIC/ST.03.07/za05to35/Train_sample/1.Calibrated/GA_M1*root"
        mars_file_list = glob.glob(
            mars_file_mask
        )  #  Here makes array which contains files matching the input condition (GA_M1*root).
        file_list = list(filter(lambda name: "123" in name, mars_file_list))
        print("Number of files : %s" % len(file_list))
    else:
        """DATA FILES"""
        mars_file_mask = (
            "/home/iwsatlas1/damgreen/CTA/MAGIC_Cleaning/datafile/Calib/*.root"
        )
        mars_file_list = glob.glob(
            mars_file_mask
        )  #  Here makes array which contains files matching the input condition (GA_M1*root).
        file_list = list(filter(lambda name: "0" in name, mars_file_list))
        print("Number of files : %s" % len(file_list))

    tolerance = 0.5

    offset = -1
    for ix, file_name in enumerate(file_list):
        magic_event_source = MAGICEventSource(input_url=file_name)
        event_generator = magic_event_source._pedestal_event_generator(telescope="M1")

        for i, event in enumerate(event_generator):
            # This is used to make sure the star and ctapipe events are lined up
            if i == 0:
                continue

            if i == 9399:
                offset -= 1
                continue

            if i > max_events:
                break

            # Basic ctapipe image processing and cleaning

            badrmspixel_mask = badpixel_calculator.get_badrmspixel_mask(event)
            deadpixel_mask = badpixel_calculator.get_deadpixel_mask(event)
            unsuitable_mask = np.logical_or(badrmspixel_mask[0], deadpixel_mask[0])

            event_image = event.dl1.tel[tel_id].image
            event_pulse_time = event.dl1.tel[tel_id].pulse_time

            clean_mask, event_image, event_pulse_time = magic_clean.clean_image(
                event_image, event_pulse_time, unsuitable_mask=unsuitable_mask
            )

            event_image_cleaned = event_image.copy()
            event_image_cleaned[~clean_mask] = 0

            try:
                hillas_params = hillas_parameters(mars_camera, event_image_cleaned)
            except:
                continue

            # All of the comparisons between star and ctapipe
            event_image_stardump = event_images_stardump[i + offset, :1039]
            event_time_stardump = event_times_stardump[i + offset, :1039]

            image_diff = (
                100 * np.abs(event_image_stardump - event_image) / event_image_stardump
            )
            image_diff_selection = image_diff > tolerance

            time_diff = (
                100
                * np.abs(event_time_stardump - event_pulse_time)
                / event_time_stardump
            )
            time_diff_selection = time_diff > tolerance

            ctapipe_unsuitable_mask = copy.copy(unsuitable_mask)
            star_unsuitable_mask = copy.copy(
                event_unsuitables_stardump[i + offset, :1039]
            )
            ctapipe_unsuitable_pixels = np.where(ctapipe_unsuitable_mask)[0]
            star_unsuitable_pixels = np.where(star_unsuitable_mask)[0]

            ctapipe_unmapped_mask = copy.copy(magic_clean.unmapped_mask)
            star_unmapped_mask = copy.copy(event_unmappeds_stardump[i + offset, :1039])
            ctapipe_unmapped_pixels = np.where(ctapipe_unmapped_mask)[0]
            star_unmapped_pixels = np.where(star_unmapped_mask)[0]

            diff_unsuitable = np.setdiff1d(
                np.union1d(ctapipe_unsuitable_pixels, star_unsuitable_pixels),
                np.intersect1d(ctapipe_unsuitable_pixels, star_unsuitable_pixels),
            )
            diff_unmapped = np.setdiff1d(
                np.union1d(ctapipe_unmapped_pixels, star_unmapped_pixels),
                np.intersect1d(ctapipe_unmapped_pixels, star_unmapped_pixels),
            )

            ctapipe_length = hillas_params.length.to("mm") / AberrationCorrection
            ctapipe_width = hillas_params.width.to("mm") / AberrationCorrection
            ctapipe_size = hillas_params.intensity

            star_idx = np.where(evtnum_arr == event.dl0.event_id)[0]
            if len(star_idx) != 1:
                continue

            star_length = length_arr[star_idx[0]] * u.mm
            star_width = width_arr[star_idx[0]] * u.mm
            star_size = size_arr[star_idx[0]]

            if star_length.value < 1e-1 or star_width.value < 1e-1:
                continue

            length_ratio = np.fabs((ctapipe_length - star_length) / star_length) * 100
            width_ratio = np.fabs((ctapipe_width - star_width) / star_width) * 100
            size_ratio = np.fabs((ctapipe_size - star_size) / star_size) * 100

            ctapipe_step1 = np.where(magic_clean.mask_step1)[0]
            star_step1 = np.asarray(event_cleanstep1[i + offset])

            ctapipe_step2 = np.where(magic_clean.mask_step2)[0]
            star_step2 = np.asarray(event_cleanstep2[i + offset])

            ctapipe_step3 = np.where(magic_clean.mask_step3)[0]
            star_step3 = np.asarray(event_cleanstep3[i + offset])

            diff_step1 = np.setdiff1d(
                np.union1d(ctapipe_step1, star_step1),
                np.intersect1d(ctapipe_step1, star_step1),
            )
            diff_step2 = np.setdiff1d(
                np.union1d(ctapipe_step2, star_step2),
                np.intersect1d(ctapipe_step2, star_step2),
            )
            diff_step3 = np.setdiff1d(
                np.union1d(ctapipe_step3, star_step3),
                np.intersect1d(ctapipe_step3, star_step3),
            )

            if (
                size_ratio > tolerance
                or length_ratio > tolerance
                or width_ratio > tolerance
            ):
                # if len(diff_step3) > 0 or len(diff_unsuitable) > 0 or len(diff_unmapped) > 0:# or np.sum(image_diff_selection) > 0 or np.sum(time_diff_selection) > 0:
                print("*" * 50)
                print(
                    bcolors.OKGREEN
                    + "Count : %i, Stereo Event Number : %s" % (i, event.dl0.event_id)
                    + bcolors.ENDC
                )

                if len(diff_unsuitable) > 0:
                    print("+" * 50)
                    print(bcolors.FAIL + "Unsuitable Diff : " + bcolors.ENDC)
                    print("diff:\t", diff_unsuitable.tolist())

                if len(diff_unmapped) > 0:
                    print("+" * 50)
                    print(bcolors.FAIL + "Unmapped Diff : " + bcolors.ENDC)
                    print("diff:\t", diff_unmapped.tolist())

                if np.sum(image_diff_selection) > 0:
                    print("+" * 50)
                    print(bcolors.FAIL + "Image Diff : " + bcolors.ENDC)
                    print(
                        "image ctapipe : ", event_image[image_diff_selection].tolist()
                    )
                    print(
                        "image dump : ",
                        event_image_stardump[image_diff_selection].tolist(),
                    )
                    print(
                        "image diff (%) : ", image_diff[image_diff_selection].tolist()
                    )

                if np.sum(time_diff_selection) > 0:
                    print("+" * 50)
                    print(bcolors.FAIL + "Time Diff : " + bcolors.ENDC)
                    print(
                        "time ctapipe : ",
                        event_pulse_time[time_diff_selection].tolist(),
                    )
                    print(
                        "time dump : ",
                        event_time_stardump[time_diff_selection].tolist(),
                    )
                    print("time diff : ", time_diff[time_diff_selection].tolist())

                string = ""
                if len(diff_step1) > 0:
                    string += "STEP1 DIFF Size %s, " % len(diff_step1)
                if len(diff_step2) > 0:
                    string += "STEP2 DIFF Size %s, " % len(diff_step2)
                if len(diff_step3) > 0:
                    string += "STEP3 DIFF Size %s" % len(diff_step3)

                if len(string) > 0:
                    print(bcolors.FAIL + string + bcolors.ENDC)

                string = ""
                if size_ratio > tolerance:
                    string += "SIZE RATIO : %.2f, " % size_ratio
                if length_ratio > tolerance:
                    string += " LENGTH RATIO : %.2f, " % length_ratio
                if width_ratio > tolerance:
                    string += " WIDTH RATIO : %.2f" % width_ratio

                if len(string) > 0:
                    print(bcolors.WARNING + string + bcolors.ENDC)
                    t = Texttable()
                    t.add_rows(
                        [
                            ["var", "ctapipe", "star", "ratio"],
                            [
                                "Size",
                                hillas_params.intensity,
                                size_arr[star_idx[0]],
                                size_ratio,
                            ],
                            [
                                "Width",
                                ctapipe_width.value,
                                star_width.value,
                                width_ratio,
                            ],
                            [
                                "Length",
                                ctapipe_length.value,
                                star_length.value,
                                length_ratio,
                            ],
                            [
                                "CorePix",
                                magic_clean.core_pix,
                                core_arr[star_idx[0]],
                                np.nan,
                            ],
                            [
                                "UsedPix",
                                magic_clean.used_pix,
                                used_arr[star_idx[0]],
                                np.nan,
                            ],
                        ]
                    )
                    print(bcolors.BOLD + bcolors.OKBLUE + t.draw() + bcolors.ENDC)

                # print("cta mask :\t" , np.where(clean_mask)[0].tolist())
                # print("star mask :\t" , star_step3.tolist())

                if len(diff_step1) > 0:
                    print(bcolors.FAIL + "Clean Step 1 Diff : " + bcolors.ENDC, end="")
                    # print("star:\t", star_step1.tolist())
                    # print("cta:\t", ctapipe_step1.tolist())
                    if len(star_step1) > len(ctapipe_step1):
                        print("Pixel not found in ctapipe : \t", diff_step1.tolist())
                    else:
                        print("Extra pixel found in ctapipe : \t", diff_step1.tolist())

                if len(diff_step2) > 0:
                    print(bcolors.FAIL + "Clean Step 2 Diff : " + bcolors.ENDC, end="")
                    # print("star:\t", star_step2.tolist())
                    # print("cta:\t", ctapipe_step2.tolist())
                    if len(star_step2) > len(ctapipe_step2):
                        print("Pixel not found in ctapipe : \t", diff_step2.tolist())
                    else:
                        print("Extra pixel found in ctapipe : \t", diff_step2.tolist())

                if len(diff_step3) > 0:
                    print(bcolors.FAIL + "Clean Step 3 Diff : " + bcolors.ENDC, end="")
                    # print("star:\t", star_step3.tolist())
                    # print("cta:\t", ctapipe_step3.tolist())
                    if len(star_step3) > len(ctapipe_step3):
                        print("Pixel not found in ctapipe : \t", diff_step3.tolist())
                    else:
                        print("Extra pixel found in ctapipe : \t", diff_step3.tolist())

            if star_size > size_cut:
                star_size_arr.append(star_size)
                ctapipe_size_arr.append(ctapipe_size)

                star_length_arr.append(star_length.value)
                ctapipe_length_arr.append(ctapipe_length.value)

                star_width_arr.append(star_width.value)
                ctapipe_width_arr.append(ctapipe_width.value)

    star_size_arr = np.asarray(star_size_arr)
    star_length_arr = np.asarray(star_length_arr)
    star_width_arr = np.asarray(star_width_arr)

    ctapipe_size_arr = np.asarray(ctapipe_size_arr)
    ctapipe_length_arr = np.asarray(ctapipe_length_arr)
    ctapipe_width_arr = np.asarray(ctapipe_width_arr)

    plt.clf()
    fig, axs = plt.subplots(2, 3, figsize=(15, 10))
    bins = 101

    axs[0, 0].scatter(np.log10(star_size_arr), np.log10(ctapipe_size_arr))
    axs[0, 0].set_title("Size")
    axs[0, 0].set_xlim([1.0, 5.5])
    axs[0, 0].set_ylim([1.0, 5.5])
    axs[0, 0].set_xlabel("Mars - log10(size)")
    axs[0, 0].set_ylabel("ctapipe - log10(size)")

    axs[0, 1].scatter(np.log10(star_length_arr), np.log10(ctapipe_length_arr))
    axs[0, 1].set_title("Length")
    axs[0, 1].set_xlim([0.0, 3.0])
    axs[0, 1].set_ylim([0.0, 3.0])
    axs[0, 1].set_xlabel("Mars - log10(length)")
    axs[0, 1].set_ylabel("ctapipe - log10(length)")

    axs[0, 2].scatter(np.log10(star_width_arr), np.log10(ctapipe_width_arr))
    axs[0, 2].set_title("Width")
    axs[0, 2].set_xlim([0.0, 3.0])
    axs[0, 2].set_ylim([0.0, 3.0])
    axs[0, 2].set_xlabel("Mars - log10(width)")
    axs[0, 2].set_ylabel("ctapipe - log10(width)")

    size_ratio = (ctapipe_size_arr - star_size_arr) / star_size_arr
    axs[1, 0].hist(size_ratio, bins=bins)
    plt.text(
        0.15,
        0.9,
        "mean : %.4f" % np.nanmean(size_ratio),
        horizontalalignment="left",
        verticalalignment="center",
        transform=axs[1, 0].transAxes,
    )
    plt.text(
        0.15,
        0.85,
        "rms : %.4f" % np.nanstd(size_ratio),
        horizontalalignment="left",
        verticalalignment="center",
        transform=axs[1, 0].transAxes,
    )
    axs[1, 0].set_xlabel("Size Ratio [(ctapipe - star)/star]")
    axs[1, 0].set_yscale("log", nonposy="clip")

    length_ratio = (ctapipe_length_arr - star_length_arr) / star_length_arr
    axs[1, 1].hist(length_ratio, bins=bins)
    plt.text(
        0.15,
        0.9,
        "mean : %.4f" % np.nanmean(length_ratio),
        horizontalalignment="left",
        verticalalignment="center",
        transform=axs[1, 1].transAxes,
    )
    plt.text(
        0.15,
        0.85,
        "rms : %.4f" % np.nanstd(length_ratio),
        horizontalalignment="left",
        verticalalignment="center",
        transform=axs[1, 1].transAxes,
    )
    axs[1, 1].set_xlabel("Length Ratio [(ctapipe - star)/star]")
    axs[1, 1].set_yscale("log", nonposy="clip")

    width_ratio = (ctapipe_width_arr - star_width_arr) / star_width_arr
    axs[1, 2].hist(width_ratio, bins=bins)
    plt.text(
        0.15,
        0.9,
        "mean : %.4f" % np.nanmean(width_ratio),
        horizontalalignment="left",
        verticalalignment="center",
        transform=axs[1, 2].transAxes,
    )
    plt.text(
        0.15,
        0.85,
        "rms : %.4f" % np.nanstd(width_ratio),
        horizontalalignment="left",
        verticalalignment="center",
        transform=axs[1, 2].transAxes,
    )
    axs[1, 2].set_xlabel("Width Ratio [(ctapipe - star)/star]")
    axs[1, 2].set_yscale("log", nonposy="clip")

    if args.base:
        plt.savefig("comp_M1_base.pdf")
    elif args.nopix:
        plt.savefig("comp_M1_base_nopix.pdf")


if __name__ == "__main__":
    main()
