from __future__ import annotations

import logging
from pathlib import Path
from typing import List

import numpy as np
import typer
from PIL import Image

from yarrtist.cli.globals import CONTEXT_SETTINGS, OPTIONS, LogLevel
from yarrtist.plotter.core_functions import (
    plot_fe_config,
    plot_fe_histo,
    plot_module_histo,
)
from yarrtist.utils.utils import (
    create_img_grid,
    data_from_config,
    fill_pdf_summary,
    format_broken_data,
    format_data,
    get_chip_type_from_config,
    get_chip_types_configs_from_scanlog,
    get_configs_from_connectivity,
    get_geomId_sn_from_config,
    get_test_list_from_scandir,
    load_data,
)

logging.basicConfig(format="[%(levelname)s] %(message)s")
log = logging.getLogger("YARRtist")

app = typer.Typer(context_settings=CONTEXT_SETTINGS)


################## plot module chip by chip ######################


def summary_perchip(
    scan_dir, tests_names, chip_names, summary_title, individual_pngs, failing
):
    all_imgs = []
    for c in chip_names:
        log.info(f"Plotting for FE {c}")
        for t in tests_names:
            try:
                log.debug(f"Plotting {t}")
                data = load_data(f"{scan_dir}/{c}_{t}.json")
            except Exception:
                log.info("Test not found for this chip, skipping")
                continue
            try:
                plotter = plot_fe_histo(data, c)
                if individual_pngs:
                    save_path = Path(f"{scan_dir}/all_yarrtist_plots/{c}/{c}_{t}.png")
                    save_path.parent.mkdir(parents=True, exist_ok=True)
                    plotter.savefig(save_path)
                fill_pdf_summary(plotter, all_imgs)
                if failing and data.get("Name") in [
                    "NoiseMap-0",
                    "ThresholdMap-0",
                    "OccupancyMap",
                    "NoiseOccupancy",
                    "Occupancy",
                ]:
                    data["Name"] = data["Name"] + " - Failures"
                    plotter = plot_fe_histo(data, c, True)
                    if individual_pngs:
                        save_path = Path(
                            f"{scan_dir}/all_yarrtist_plots/{c}/{c}_{t}_highlighted.png"
                        )
                        plotter.savefig(save_path)
                    fill_pdf_summary(plotter, all_imgs)
            except Exception as e:
                log.warning(f"Failed plotting {t}: {e}")

    if all_imgs:
        summary = [Image.open(p).convert("RGB") for p in all_imgs]

        summary[0].save(
            f"{scan_dir}/{summary_title}.pdf",
            save_all=True,
            append_images=summary[1:],
        )
        log.info(f"Plot summary saved in {scan_dir}/{summary_title}.pdf")
        if individual_pngs:
            log.info(
                f"Individual pngs for each chip saved in {scan_dir}/all_yarrtist_plots"
            )
    else:
        log.info("Plot summary is empty")


def config_summary_perchip(scan_dir, chip_names, chip_types, summary_title):
    all_imgs = []
    files_list = [f.name for f in Path(scan_dir).iterdir()]

    before_configs = [
        load_data(f"{scan_dir}/{f}") for f in files_list if ".before" in f
    ]
    after_configs = [load_data(f"{scan_dir}/{f}") for f in files_list if ".after" in f]

    for chip, type in zip(chip_names, chip_types):
        if type not in ["RD53B", "ITKPIXV2"]:
            log.warning(
                f"{chip}: config plotting for {type} not available at the moment"
            )
            continue
        log.info(f"Plotting for FE {chip}")
        configs = {
            "before": [
                data
                for data in before_configs
                if data[get_chip_type_from_config(data)]["Parameter"]["Name"] == chip
            ],
            "after": [
                data
                for data in after_configs
                if data[get_chip_type_from_config(data)]["Parameter"]["Name"] == chip
            ],
        }
        tdac_distr = {}
        if len(configs["before"]) == 0 and len(configs["after"]) == 0:
            continue
        for ba, config in configs.items():
            if len(config) == 0:
                continue
            images = []
            config_data = config[0]
            try:
                plotters = plot_fe_config(config_data, f"{chip} ({ba})")
                for k in plotters:
                    if k != "TDAC_1d":
                        fill_pdf_summary(plotters[k], images)
                chip_img = create_img_grid(images, (2, 2), f"{chip} ({ba})")
                all_imgs.append(chip_img)
                tdac_distr[f"{ba}"] = plotters["TDAC_1d"]
            except Exception as e:
                log.warning(f"Failed to plot: {e}")
        try:
            tdac_imgs = []
            for _k, p in tdac_distr.items():
                fill_pdf_summary(p, tdac_imgs)
            chip_img = create_img_grid(tdac_imgs, (1, 2), "")
            all_imgs.append(chip_img)
        except Exception as e:
            log.warning(f"Failed to plot: {e}")

    if all_imgs:
        summary = [p.convert("RGB") for p in all_imgs]
        summary[0].save(
            f"{scan_dir}/{summary_title}.pdf",
            save_all=True,
            append_images=all_imgs[1:],
        )
        log.info(f"Config summary saved in {scan_dir}/{summary_title}.pdf")
    else:
        log.info("Config summary is empty")


################### summary of module from YARR scan directory ###################
def module_summary(
    module_info,
    tests_names,
    module_type,
    scan_dir,
    summary_title,
    failing,
    individual_pngs,
):
    images = []
    ctype = module_info.get("chipType")
    all_ch_names = [next(iter(chip.keys())) for chip in module_info.get("chips", [])]

    for test in tests_names:
        plot = {}
        plot["data"] = {}

        log.info(f"Plotting {test}")
        not_found = 0

        # for geomId, item in enumerate(ordered_chips):
        for id, chip in enumerate(all_ch_names):
            # log.info(chip)
            # chip = item.get("name")
            # ctype = item.get("type")
            geomId = id
            file_name = f"{chip}_{test}.json"
            log.debug(f"Plotting for chip {geomId + 1}: {chip}")
            if file_name not in [f.name for f in Path(scan_dir).iterdir()]:
                log.debug(f"Chip {chip} not found, filling with 0s")
                not_found += 1
                format_broken_data(plot, test, ctype, geomId)
                continue

            data = load_data(f"{scan_dir}/{file_name}")
            format_data(plot, data, geomId)

        if (not_found == 4 and module_type == "Quad") or (
            not_found == 3 and module_type == "Triplet"
        ):
            log.debug("Skipping whole test for this module")
            continue
        try:
            plotter = plot_module_histo(plot, data, module_type, "", all_ch_names)
            if individual_pngs:
                save_path = Path(
                    f"{scan_dir}/yarrtist_{summary_title}/{plot.get('title')}.png"
                )
                save_path.parent.mkdir(parents=True, exist_ok=True)
                plotter.savefig(save_path)
            fill_pdf_summary(plotter, images)
            if failing and plot.get("title") in [
                "NoiseMap-0",
                "ThresholdMap-0",
                "OccupancyMap",
                "NoiseOccupancy",
                "Occupancy",
            ]:
                plot["title"] = plot.get("title") + " - Failures"
                plotter = plot_module_histo(
                    plot,
                    data,
                    module_type,
                    "",
                    all_ch_names,
                    True,
                )
                if individual_pngs:
                    save_path = Path(
                        f"{scan_dir}/yarrtist_{summary_title}/{plot.get('title')}_highlighted.png"
                    )
                    plotter.savefig(save_path)
                fill_pdf_summary(plotter, images)
        except Exception as e:
            log.warning(f"Failed to plot {test}: {e}")

    if images:
        summary = [Image.open(p).convert("RGB") for p in images]
        summary[0].save(
            f"{scan_dir}/{summary_title}.pdf",
            save_all=True,
            append_images=summary[1:],
        )
        log.info(f"Plot summary saved in {scan_dir}/{summary_title}.pdf")
    else:
        log.info("Plot summary is empty")


################### summary of module config from YARR scan directory ###################
def module_config_summary(module_type, module_info, scan_dir, summary_title):
    all_data_before = []
    all_data_after = []
    all_imgs = []

    all_ch_names = [next(iter(chip.keys())) for chip in module_info.get("chips", [])]
    all_files = [f.name for f in Path(scan_dir).iterdir()]
    before_configs = [load_data(f"{scan_dir}/{f}") for f in all_files if ".before" in f]
    after_configs = [load_data(f"{scan_dir}/{f}") for f in all_files if ".after" in f]

    for chip in all_ch_names:
        before_cfg = [
            data_from_config(data)
            for data in before_configs
            if data[get_chip_type_from_config(data)]["Parameter"]["Name"] == chip
        ]
        after_cfg = [
            data_from_config(data)
            for data in after_configs
            if data[get_chip_type_from_config(data)]["Parameter"]["Name"] == chip
        ]
        if len(before_cfg) == 0:
            all_data_before.append([{}] * 5)
        else:
            all_data_before.append(before_cfg[0])
        if len(after_cfg) == 0:
            all_data_after.append([{}] * 5)
        else:
            all_data_after.append(after_cfg[0])

    all_data_before = list(map(list, zip(*all_data_before)))
    all_data_after = list(map(list, zip(*all_data_after)))

    for n in range(len(all_data_after)):
        plot_before = {}
        plot_after = {}
        plot_before["data"] = {}
        plot_after["data"] = {}
        images = []
        for geomId, (tb, ta) in enumerate(zip(all_data_before[n], all_data_after[n])):
            if not tb or not ta:
                plot_before["data"].update({geomId: np.zeros((400, 384))})
                plot_after["data"].update({geomId: np.zeros((400, 384))})
                continue
            format_data(plot_before, tb, geomId)
            plot_before["title"] = plot_before["title"] + " (before)"
            format_data(plot_after, ta, geomId)
            plot_after["title"] = plot_after["title"] + " (after)"
        try:
            plotter = plot_module_histo(plot_before, tb, module_type, "", all_ch_names)
            fill_pdf_summary(plotter, images)
            plotter = plot_module_histo(plot_after, ta, module_type, "", all_ch_names)
            fill_pdf_summary(plotter, images)
            all_imgs.append(create_img_grid(images, (1, 2)))
        except Exception as e:
            log.warning(f"Failed plotting: {e}")

    if all_imgs:
        summary = [p.convert("RGB") for p in all_imgs]

        summary[0].save(
            f"{scan_dir}/{summary_title}.pdf",
            save_all=True,
            append_images=all_imgs[1:],
        )
        log.info(f"Config summary saved in {scan_dir}/{summary_title}.pdf")
    else:
        log.info("Config summary is empty")


@app.command()
def main(
    connectivity_files: List[Path] = OPTIONS["connectivity_files"],
    scan_directory: Path = OPTIONS["scan_directory_scan_plot"],
    relative_path: Path = OPTIONS["path"],
    per_module: bool = OPTIONS["per_module"],
    config_summary: bool = OPTIONS["config_summary"],
    failing: bool = OPTIONS["failing"],
    individual_pngs: bool = OPTIONS["individual_pngs"],
    verbosity: LogLevel = OPTIONS["verbosity"],
):
    log.setLevel(verbosity.value)

    if not connectivity_files:
        all_configs = get_chip_types_configs_from_scanlog(
            load_data(f"{scan_directory}/scanLog.json"), scan_directory, relative_path
        )
    else:
        log.info("Using connectivity file to get chips configurations")
        all_configs = []
        for conn_file in connectivity_files:
            all_configs.append(
                get_configs_from_connectivity(
                    load_data(conn_file), Path(conn_file).parent, relative_path
                )
            )

    all_chip_types = []
    for m in all_configs:
        chip_type = m.get("chipType")
        all_chip_types.extend([chip_type] * len(m.get("configs")))
        m["chips"] = []
        for c in m.get("configs"):
            full_config = load_data(f"{c}")
            geomId, name = get_geomId_sn_from_config(full_config, chip_type)
            m["chips"].append({f"{name}": {"id": geomId}})  # , "config": full_config}})
        if per_module:
            m["chips"] = sorted(m["chips"], key=lambda d: next(iter(d.values()))["id"])

    all_chip_names = [
        next(iter(chip.keys())) for m in all_configs for chip in m.get("chips", [])
    ]
    all_tests = get_test_list_from_scandir(all_chip_names, scan_directory)

    if not per_module:
        log.info(f"Loaded {len(all_chip_names)} chips")
        log.info("Plotting scan summary, divided per chip")

        summary_perchip(
            scan_directory,
            all_tests,
            all_chip_names,
            "summary_perchip",
            individual_pngs,
            failing,
        )

        if config_summary:
            log.info("Plotting config summary, divided per chip")

            config_summary_perchip(
                scan_directory,
                all_chip_names,
                all_chip_types,
                "config_summary_perchip",
            )

    else:
        log.info("Plotting module scan summary")
        log.info(f"Loaded {len(all_configs)} modules")

        for i, config in enumerate(all_configs):
            log.info(f"Plotting for module {i + 1}")
            chip_type = config.get("chipType")
            if chip_type not in ["RD53B", "ITKPIXV2"]:
                log.error("This option is available only for RD53B and ITKPIXV2")
            else:
                n_chips = len(config.get("chips", []))
                if n_chips not in [3, 4]:
                    log.error("Number of chips does not match a known module type")
                else:
                    module_type = "Triplet" if n_chips == 3 else "Quad"
                    log.info(f"Module is a {module_type}")

                    module_summary(
                        config,
                        all_tests,
                        module_type,
                        scan_directory,
                        f"module_summary_{i}",
                        failing,
                        individual_pngs,
                    )

                    if config_summary:
                        log.info(f"Plotting module {i + 1} config summary")
                        module_config_summary(
                            module_type,
                            config,
                            scan_directory,
                            f"module_config_summary_{i}",
                        )
