#!/usr/bin/env python
# coding: utf-8

import argparse
import yaml
import pandas as pd
import uproot
from ctapipe_io_magic import MAGICEventSource
from ctapipe.visualization import CameraDisplay
from ctapipe.instrument import CameraGeometry
from ctapipe.image import hillas_parameters
from ctapipe.image.timing import timing_parameters
from ctapipe.instrument import CameraDescription
from utils import MAGIC_Badpixels
from utils import MAGIC_Cleaning
import matplotlib.pyplot as plt
import matplotlib
import numpy as np

from ctapipe.image.morphology import number_of_islands
from mars_images_to_hdf5 import read_images

# ----------------------------------------------------------------

# create a list of the events you want to be compared
ids_to_compare = []

# add argparser arguments
arg_parser = argparse.ArgumentParser(
    description="""
This tool compares the images of events processed by MARS and the magic-cta-pipeline. 
The output is a png file with the camera images and a hdf5 file that contains the pixel information.
"""
)

arg_parser.add_argument(
    "--config",
    default="config.yaml",
    help="Configuration file to steer the code execution.",
)
arg_parser.add_argument(
    "--useall",
    help="Compare images for all events in the input file",
    action="store_true",
)
arg_parser.add_argument(
    "--use_ids_config", help="Compare images for certain events", action="store_true"
)

parsed_args = arg_parser.parse_args()

file_not_found_message = """
Error: can not load the configuration file {:s}.
Please check that the file exists and is of YAML or JSON format.
Exiting.
"""


# load get_images_config file
try:
    config = yaml.safe_load(open(parsed_args.config, "r"))
except IOError:
    print(file_not_found_message.format(parsed_args.config))
    exit()
# check for the important keys
if "input_files" not in config:
    print(
        'Error: the configuration file is missing the "input_files" section. Exiting.'
    )
    exit()

if "output_files" not in config:
    print(
        'Error: the configuration file is missing the "output_files" section. Exiting.'
    )
    exit()


# ---------------------------------------------------------
#    Compare images for all events in the given file
# ---------------------------------------------------------
if parsed_args.useall:
    mcp_input = config["input_files"]["magic-cta-pipe"]
    tel_id = config["information"]["tel_id"]
    # does only work for stereo files
    mcp_file = uproot.open(config["input_files"]["magic-cta-pipe"])
    # print(mcp_file.keys())
    if "Events" not in mcp_file.keys():
        print("Missing key: Events. Unacceptable file format. Exiting.")
        exit()
    else:
        if (
            "MRawEvtHeader./MRawEvtHeader.fStereoEvtNumber"
            not in mcp_file["Events"].keys()
        ):
            print(
                f"Missing key: MRawEvtHeader_{tel_id}.fStereoEvtNumber. Unacceptable file format. Exiting."
            )
            exit()
        else:
            df_mcp_data = pd.DataFrame()
            df_mcp_data["event_id"] = (
                mcp_file["Events"]["MRawEvtHeader./MRawEvtHeader.fStereoEvtNumber"]
                .array()
                .to_numpy()
            )
            ids_to_compare = df_mcp_data["event_id"].tolist()
            # the keys depend on the file, there may be a case where they have to be changed
# ---------------------------------------------------------
#       Compare images for only certain events
# ---------------------------------------------------------

elif parsed_args.use_ids_config:
    ids_to_compare = config["event_list"]

else:
    print("No comparing argument specified. Exiting.")
    exit()


# print the list
if len(ids_to_compare) >= 21:
    print("EVENTS:", ids_to_compare[:10], "...", ids_to_compare[-10:])
else:
    print("EVENTS:", ids_to_compare)

print(len(ids_to_compare))


# ------------------------------------------------------------
# define camera geometry
def new_camera_geometry(camera_geom):
    """Scale given camera geometry of a given (constant) factor

    Parameters
    ----------
    camera : CameraGeometry
        Camera geometry
    factor : float
        Scale factor

    Returns
    -------
    CameraGeometry
        Scaled camera geometry
    """

    return CameraGeometry(
        camera_name="MAGICCam",
        pix_id=camera_geom.pix_id,
        pix_x=-1.0 * camera_geom.pix_y,
        pix_y=-1.0 * camera_geom.pix_x,
        pix_area=camera_geom.guess_pixel_area(
            camera_geom.pix_x, camera_geom.pix_y, camera_geom.pix_type
        ),
        pix_type=camera_geom.pix_type,
        pix_rotation=camera_geom.pix_rotation,
        cam_rotation=camera_geom.cam_rotation,
    )


# get information from the config file
telescope_id = config["information"]["tel_id"]
run_num = config["information"]["run_number"]
out_path = config["output_files"]["file_path"]


# compare events one at a time
ids = []
for id_to_compare in ids_to_compare:
    ids.append(id_to_compare)
    print("EVENT ID:", ids)
    for id_event in ids:
        # ------------------
        #       MARS
        # -----------------

        mars_input = config["input_files"]["mars"]

        image = []  # pixel charges
        events = []  # event ids
        telescope = []  # telescope id
        runnum = []  # run number

        for image_container in read_images(mars_input):
            events.append(image_container.event_id)
            telescope.append(image_container.tel_id)
            runnum.append(image_container.obs_id)
            image.append(image_container.image)

        image = image.to_numpy()
        stereo_id = events.to_numpy()

        MAGICCAM = CameraDescription.from_name("MAGICCam")
        GEOM = MAGICCAM.geometry
        geometry_mars = new_camera_geometry(GEOM)

        # uncomment this list to compare different events
        # ids = [1962]

        for id in range(len(image)):
            if stereo_id[id] not in ids:
                continue
            event_image = np.array(image[id][:1039])
            clean_mask = event_image != 0
            print("Event ID:", stereo_id[id])
            print(event_image[clean_mask])
            # print(event_image)
            event_image_mars = event_image.copy()

            # compare number of islands
            num_islands_mars, island_labels_mars = number_of_islands(
                geometry_mars, (np.array(event_image[:1039]) > 0)
            )
            # print(num_islands_mars)

        # -----------------------------
        #        magic-cta-pipe
        # -----------------------------

        source = MAGICEventSource(input_url=config["input_files"]["mcp_source"])

        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=False,
        )

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

        # uncomment this list to compare different events
        # ids = [1962]

        for event in source._mono_event_generator(telescope=f"{telescope_id}"):
            if event.index.event_id not in ids:
                continue
            print(f"Event ID: {event.index.event_id}")
            tel = source.subarray.tel
            for tel_id in [telescope_id]:
                print("Telescope ID:", tel_id)
                r0tel = event.r0.tel[tel_id]
                geometry_old = tel[tel_id].camera.geometry
                geometry_mcp = new_camera_geometry(geometry_old)
                magic_clean = MAGIC_Cleaning.magic_clean(geometry_mcp, cleaning_config)
                badpixel_calculator = MAGIC_Badpixels.MAGICBadPixelsCalc(
                    config=bad_pixels_config
                )
                event_image = event.dl1.tel[tel_id].image
                event_pulse_time = event.dl1.tel[tel_id].peak_time

                """ 
                find different catgories of pixels
                ------
                bad: pixel that is unsuitable, defined subrunwise
                hot: too high charge due to star or not working correctly
                badrms: the pedestal rms value is over a certain threshhold; this is bad and hot pixels combined
                    get badrms pixels via badrmspixel_mask
                    get badrms pixel indices via bad_pixel_indices
                dead: this pixels is not working at all
                    get dead pixels via deadpixel_mask 
                    get dead pixel indices via dead_pixel_indices 
                unsuitable: union of badrms and dead pixels
                    get unsuitable pixels via unsuitable_mask
                ---------
                """

                badrmspixel_indices = [[None], [None]]

                badrmspixel_mask = badpixel_calculator.get_badrmspixel_mask(event)

                deadpixel_mask = badpixel_calculator.get_deadpixel_mask(event)
                unsuitable_mask = np.logical_or(
                    badrmspixel_mask[tel_id - 1], deadpixel_mask[tel_id - 1]
                )

                bad_pixel_indices = [
                    i for i, x in enumerate(badrmspixel_mask[telescope_id - 1]) if x
                ]
                # print("badpixel_indices:", bad_pixel_indices)
                dead_pixel_indices = [
                    i for i, x in enumerate(deadpixel_mask[telescope_id - 1]) if x
                ]
                # print("deadpixel_indices:", dead_pixel_indices)
                bad_not_dead_pixels = []
                for i in bad_pixel_indices:
                    if i not in dead_pixel_indices:
                        bad_not_dead_pixels.append(i)
                # print("bad but not dead pixels:", bad_not_dead_pixels)
                # print(len(bad_pixel_indices))

                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
                # print(event_image_cleaned[clean_mask])

                event_pulse_time_cleaned = event_pulse_time.copy()
                event_pulse_time_cleaned[~clean_mask] = 0

                if np.any(event_image_cleaned):
                    hillas_params = hillas_parameters(geometry_mcp, event_image_cleaned)
                    image_mask = event_image_cleaned > 0
                    timing_params = timing_parameters(
                        geometry_mcp,
                        event_image_cleaned,
                        event_pulse_time_cleaned,
                        hillas_params,
                        image_mask,
                    )
                    print(timing_params.slope)
                else:
                    print("Cleaning failed.")
                    continue

                # compare number of islands
                num_islands_mcp, island_labels_mcp = number_of_islands(
                    geometry_mcp, (np.array(event_image_cleaned[:1039]) > 0)
                )
                # print(num_islands_mcp)

        # ------------------
        # original calibrated data
        # ---------------------
        # chooses the path for either M1 or M2
        # specified by the telescope id in the config file
        data_path = config["input_files"][f"original_data_M{telescope_id}"]
        with uproot.open(data_path) as input_data:
            event_ids = input_data["Events"]["MRawEvtHeader.fStereoEvtNumber"].array()
            images = input_data["Events"]["MCerPhotEvt.fPixels.fPhot"].array()

            event_index_array = np.where(event_ids == id_event)
            event_index = event_index_array[0][0]
            # print(event_index)
            calibrated_data_images = images[event_index][:1039]
            # print(np.array(images[event_index][:1039]).shape)
            # print(np.array(images).shape)
            # print(np.array(calibrated_data_images).shape)

            # print(len(data_images))

        # -------------------------------------------
        # get maximum pixel charge for the colorbar
        # -------------------------------------------

        mcp_max = np.amax(event_image_cleaned[clean_mask])
        mars_max = np.amax(event_image[clean_mask])
        if mcp_max >= mars_max:
            vmax = mcp_max
        else:
            vmax = mars_max

        vmin = 0

        # ---------------------------------------------
        # find pixel differences and write h5 output file
        # -------------------------------------------------

        pix_mars = np.array([event_image_mars])
        pix_mcp = np.array([event_image_cleaned])
        pix_diff = abs(pix_mars - pix_mcp)
        # if you want negative values instead of the absolutes
        # pix_diff = pix_mars - pix_mcp
        # pix_diff = abs(pix_mars - pix_mcp[0])
        pix_image = pix_diff[0]
        clean_mask_pixels = pix_image != 0

        # differences between outputs from MARS and mcp and input calibrated data
        pix_diff_mars = abs(pix_mars[0] - calibrated_data_images)
        pix_diff_mcp = abs(pix_mcp[0] - calibrated_data_images)
        print(pix_diff_mars)
        # print(calibrated_data_images)

        data_value = []

        for i in range(1039):
            data_value.append(calibrated_data_images[i])

        df_pixel = pd.DataFrame()
        data = np.array([event_image_mars, event_image_cleaned])
        data_real = np.transpose(data)

        df_pixel = pd.DataFrame(data_real, columns=["MARS charge", "mcp charge"])
        df_pixel["calibrated data"] = np.transpose(
            data_value
        )  # (calibrated_data_images[:1039])
        df_pixel["difference MARS mcp"] = np.transpose(pix_diff)
        df_pix_diff = df_pixel.loc[df_pixel["difference MARS mcp"] > 0]
        pix_diff_ids = df_pix_diff.index.tolist()
        print(df_pix_diff)

        # saving the output
        if config["save_only_when_differences"] == True:
            # the file only gets saved if there are differences between the images
            if any(pix_image) == True:
                print("Differences found. Saving files!")
                df_pixel.to_hdf(
                    f"{out_path}{run_num}_{id_event}_M{telescope_id}_pixel_info.h5",
                    "/pixel_information",
                    "w",
                )
                df_pix_diff.to_hdf(
                    f"{out_path}{run_num}_{id_event}_M{telescope_id}_pixel_diff.h5",
                    "/pixel_differences",
                    "w",
                )
                print(pix_diff_ids)

            else:
                print("No differences found. No files will be saved!")
                continue

        elif config["save_only_when_differences"] == False:
            # the file gets saved in any case
            df_pixel.to_hdf(
                f"{out_path}{run_num}_{id_event}_M{telescope_id}_pixel_info.h5",
                "/pixel_information",
                "w",
            )
            df_pix_diff.to_hdf(
                f"{out_path}{run_num}_{id_event}_M{telescope_id}_pixel_diff.h5",
                "/pixel_differences",
                "w",
            )
            print(pix_diff_ids)
        else:
            print("No criteria for saving data specified. Exiting.")
            exit()

        # --------------------------------
        #  creating output image
        # ----------------------------
        fig = plt.figure(figsize=(20, 10))
        grid_shape = (2, 3)

        # original data
        plt.subplot2grid(grid_shape, (0, 0))

        geom = CameraGeometry.from_name("MAGICCamMars")
        disp = CameraDisplay(geom, calibrated_data_images)
        # pixels whose original value is negative
        negative_mask = calibrated_data_images < 0
        disp.highlight_pixels(negative_mask, color="white", alpha=1, linewidth=1)
        disp.add_colorbar(label="pixel charge")
        disp.set_limits_minmax(vmin, vmax)

        plt.title("original data")

        # mars_data
        plt.subplot2grid(grid_shape, (0, 1))

        norm = matplotlib.colors.Normalize(vmin=0, vmax=vmax)
        disp_mars = CameraDisplay(geometry_mars, event_image_mars[:1039])
        # disp_mars.highlight_pixels(clean_mask, color='white', alpha=0.5, linewidth=1)
        disp_mars.set_limits_minmax(vmin, vmax)
        disp_mars.add_colorbar(label="pixel charge")
        disp_mars.highlight_pixels(clean_mask_pixels, color="red", alpha=1, linewidth=1)

        plt.title("MARS data")

        # mcp_data
        plt.subplot2grid(grid_shape, (0, 2))

        disp_mcp = CameraDisplay(geometry_mcp, image=event_image_cleaned)
        # disp_mcp.highlight_pixels(clean_mask, color='white', alpha=0.5, linewidth=1)
        disp_mcp.add_colorbar(label="pixel charge")
        disp_mcp.set_limits_minmax(vmin, vmax)
        disp_mcp.highlight_pixels(clean_mask_pixels, color="red", alpha=1, linewidth=1)

        plt.title("magic_cta_pipe data")

        # differences between MARS and mcp
        plt.subplot2grid(grid_shape, (1, 0))

        disp = CameraDisplay(geom, pix_image[:1039])
        disp.add_colorbar(label="pixel charge")
        disp.highlight_pixels(clean_mask_pixels, color="red", alpha=1, linewidth=1)

        # alternative image: used pixels after the cleaning and differences mask
        """for id in range(len(image)):
           if stereo_id[id] not in ids:
                continue
            event_image = np.array(image[id][:1039])
            clean_mask = event_image != 0
            #print(event_image[clean_mask])
            norm = matplotlib.colors.Normalize(vmin=0, vmax=vmax)
            disp = CameraDisplay(geometry_mars, image=event_image)
            disp.highlight_pixels(clean_mask_pixels, color='red', alpha=0.5, linewidth=2)
            disp.set_limits_minmax(vmin, 1)"""

        plt.title("differences MARS-mcp")

        # differences between MARS and the calibrated data
        plt.subplot2grid(grid_shape, (1, 1))

        pix_diff_mars_copy = np.array(pix_diff_mars).copy()
        pix_diff_mars_copy[np.array(event_image_mars[:1039]) == 0] = 0

        disp = CameraDisplay(geom, pix_diff_mars_copy > 0)
        # disp.highlight_pixels(clean_mask_pixels, color='red', alpha=1, linewidth=1)
        disp.highlight_pixels(
            np.array(event_image_mars[:1039]) != 0, color="white", alpha=1, linewidth=3
        )
        # disp.set_limits_minmax(vmin, vmax)
        # disp.add_colorbar(label='pixel charge')

        plt.title("differences MARS-original data")

        # mcp-orig-data-diff
        plt.subplot2grid(grid_shape, (1, 2))

        pix_diff_mcp_copy = np.array(pix_diff_mcp).copy()
        pix_diff_mcp_copy[np.array(event_image_cleaned) == 0] = 0

        disp = CameraDisplay(geom, pix_diff_mcp_copy > 0)
        disp.highlight_pixels(
            np.array(event_image_cleaned) != 0, color="white", alpha=1, linewidth=3
        )
        # disp.set_limits_minmax(vmin, vmax)
        # disp.add_colorbar(label='pixel charge')
        plt.title("differences mcp-original data")

        fig.suptitle(
            f"Comparison_MARS_magic-cta-pipe: Event ID {id_event}, {run_num}, M{tel_id}",
            fontsize=16,
        )
        fig.savefig(f"{out_path}/data-comparison-{run_num}_{id_event}_M{tel_id}.png")
        # fig.savefig(f"{out_path}/data-comparison-{run_num}_{id_event}_M{tel_id}.pdf")

    ids.clear()
