import atexit
import csv
import shutil
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any

import jax
import jax.numpy as jnp
import numpy as np
import seaborn as sns
from loguru import logger
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from rich.console import Console
from rich.markup import escape
from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
from rich.table import Table

from fdtdx.core.conversion.export import export_stl as export_stl_fn
from fdtdx.core.misc import cast_floating_to_numpy
from fdtdx.core.plotting.device_permittivity_index_utils import device_matrix_index_figure
from fdtdx.objects.container import ObjectContainer, ParameterContainer
from fdtdx.objects.detectors.detector import DetectorState


def init_working_directory(experiment_name: str, wd_name: str | None) -> Path:
    """Initialize working directory for experiment outputs.

    Creates a timestamped directory structure for experiment outputs under outputs/nobackup/.
    Uses current date/time unless a specific working directory name is provided.

    Args:
        experiment_name: Name of the experiment
        wd_name: Optional specific name for the working directory. If None, uses timestamp.

    Returns:
        Path: Created working directory path
    """
    cur_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S.%f")
    day, daytime = cur_time.split("_")
    new_cwd = (
        Path(__file__).parents[2]
        / "outputs"
        / "nobackup"
        / day
        / experiment_name
        / (daytime if wd_name is None else wd_name)
    )
    new_cwd.mkdir(parents=True)
    return new_cwd


def _log_formatter(record: Any) -> str:
    """Log message formatter"""
    color_map = {
        "TRACE": "dim blue",
        "DEBUG": "cyan",
        "INFO": "bold",
        "SUCCESS": "bold green",
        "WARNING": "yellow",
        "ERROR": "bold red",
        "CRITICAL": "bold white on red",
    }
    lvl_color = color_map.get(record["level"].name, "cyan")
    loc = record["file"].path + ":" + str(record["line"])
    message = escape(record["message"])
    return (
        "[not bold green]{time:DD.MM.YYYY HH:mm:ss.SSS}[/not bold green] | "
        + f"{loc}"
        + f" - [{lvl_color}]{message}[/{lvl_color}]"
    )


def snapshot_python_files(snapshot_dir: Path):
    snapshot_dir.mkdir(parents=True, exist_ok=True)
    # fdtdx
    root_dir = Path(__file__).parent.parent
    files = list(root_dir.rglob("*.py"))
    # scripts
    scripts_dir = Path(__file__).parent.parent.parent / "scripts"
    files = files + list(scripts_dir.rglob("*.py"))

    for python_file in files:
        relative_path = python_file.relative_to(root_dir.parent)
        destination = snapshot_dir / relative_path
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(python_file, destination)
    # make zip and delete directory
    shutil.make_archive(str(snapshot_dir.parent / "code"), "zip", snapshot_dir)
    shutil.rmtree(snapshot_dir)


class Logger:
    """Logger for managing experiment outputs and visualization.

    Handles experiment logging, metrics tracking, and visualization of simulation results.
    Creates a working directory structure, initializes logging, and provides methods for
    saving figures, metrics, and device parameters.

    Args:
        experiment_name: Name of the experiment
        name: Optional specific name for the working directory. If None, uses timestamp.
    """

    def __init__(self, experiment_name: str, name: str | None = None):
        sns.set_theme(context="paper", style="white", palette="colorblind")
        self.cwd = init_working_directory(experiment_name, wd_name=name)
        self.console = Console()
        self.progress = Progress(
            SpinnerColumn(),
            *Progress.get_default_columns(),
            TimeElapsedColumn(),
            console=self.console,
        ).__enter__()
        atexit.register(self.progress.stop)
        logger.remove()
        logger.add(
            self.console.print,
            level="TRACE",
            format=_log_formatter,
            colorize=True,
        )
        logger.add(
            self.cwd / "logs.log",
            level="TRACE",
            format="{time:DD.MM.YYYY HH:mm:ss:ssss} | {level} - {message}",
        )
        logger.info(f"Starting experiment {experiment_name} in {self.cwd}")
        snapshot_python_files(self.cwd / "code")
        self.fieldnames = None
        self.writer = None
        self.csvfile = open(self.cwd / "metrics.csv", "w", newline="")
        self.last_indices: dict[str, jax.Array | None] = defaultdict(lambda: None)
        atexit.register(self.csvfile.close)

    @property
    def stl_dir(self) -> Path:
        """Directory for storing STL files.

        Returns:
            Path: Directory for STL file outputs
        """
        directory = self.cwd / "device" / "stl"
        directory.mkdir(parents=True, exist_ok=True)
        return directory

    @property
    def params_dir(self) -> Path:
        """Directory for storing parameter files.

        Returns:
            Path: Directory for parameter file outputs
        """
        directory = self.cwd / "device" / "params"
        directory.mkdir(parents=True, exist_ok=True)
        return directory

    def savefig(self, directory: Path, filename: str, fig: Figure, dpi: int = 300):
        """Save a matplotlib figure to file.

        Creates a figures subdirectory if needed and saves the figure with specified settings.

        Args:
            directory: Base directory to save in
            filename: Name for the figure file
            fig: Matplotlib figure to save
            dpi: Resolution in dots per inch
        """
        figure_directory = directory / "figures"
        figure_directory.mkdir(parents=True, exist_ok=True)
        fig.savefig(directory / "figures" / filename, dpi=dpi, bbox_inches="tight")
        plt.close(fig)

    def write(self, stats: dict, do_print: bool = True):
        """Write statistics to CSV file and optionally print them.

        Records metrics in a CSV file and optionally displays them in a formatted table.
        Automatically initializes CSV headers on first write.

        Args:
            stats: Dictionary of statistics to record
            do_print: Whether to print stats to console
        """
        stats = {
            k: v.item() if isinstance(v, jax.Array) else v
            for k, v in stats.items()
            if isinstance(v, (int, float)) or (isinstance(v, jax.Array) and v.size == 1)
        }
        if self.fieldnames is None:
            self.fieldnames = list(stats.keys())
            self.writer = csv.DictWriter(self.csvfile, fieldnames=self.fieldnames)
            self.writer.writeheader()
        assert self.writer is not None
        self.writer.writerow(stats)
        self.csvfile.flush()
        if do_print:
            table = Table(box=None)
            for k, v in stats.items():
                table.add_column(k)
                table.add_column(str(v))
            self.console.print(table)

    def log_detectors(
        self,
        iter_idx: int,
        objects: ObjectContainer,
        detector_states: dict[str, DetectorState],
        exclude: list[str] = [],
    ):
        """Log detector states and generate visualization plots.

        Creates plots for each detector's state and saves them to the detector's output directory.
        Handles both figure outputs and other detector-specific file formats.

        Args:
            iter_idx: Current iteration index
            objects: Container with simulation objects
            detector_states: Dictionary mapping detector names to their states
            exclude: List of detector names to exclude from logging
        """
        for detector in [d for d in objects.detectors if d.name not in exclude]:
            cur_state = jax.device_get(detector_states[detector.name])
            cur_state = cast_floating_to_numpy(cur_state, float)

            if not detector.plot:
                continue
            figure_dict = detector.draw_plot(
                state=cur_state,
                progress=self.progress,
            )

            detector_dir = self.cwd / "detectors" / detector.name
            detector_dir.mkdir(parents=True, exist_ok=True)

            for k, v in figure_dict.items():
                if isinstance(v, Figure):
                    self.savefig(
                        detector_dir,
                        f"{detector.name}_{k}_{iter_idx}.png",
                        v,
                        dpi=detector.plot_dpi,  # type: ignore
                    )
                elif isinstance(v, str):
                    shutil.copy(
                        v,
                        detector_dir / f"{detector.name}_{k}_{iter_idx}{Path(v).suffix}",
                    )
                else:
                    raise Exception(f"invalid detector output for plotting: {k}, {v}")

    def log_params(
        self,
        iter_idx: int,
        params: ParameterContainer,
        objects: ObjectContainer,
        export_figure: bool = False,
        export_stl: bool = False,
        export_air_stl: bool = False,
    ) -> int:
        """Log parameter states and export device visualizations.

        Saves device parameters and optionally exports visualizations as figures or STL files.
        Tracks changes in device voxels between iterations.

        Args:
            iter_idx: Current iteration index
            params: Container with device parameters
            objects: Container with simulation objects
            export_figure: Whether to export index matrix figures
            export_stl: Whether to export device geometry as STL
            export_air_stl: Whether to export air regions as STL

        Returns:
            int: Number of voxels that changed since last iteration
        """
        changed_voxels = 0
        for device in objects.devices:
            device_params = params[device.name]
            indices = device.get_indices(device_params)

            has_previous = self.last_indices[device.name] is not None
            cur_changed_voxels = 0
            if has_previous:
                last_device_indices = self.last_indices[device.name]
                cur_changed_voxels = int(jnp.sum(indices != last_device_indices))
            changed_voxels += cur_changed_voxels
            self.last_indices[device.name] = indices
            if cur_changed_voxels == 0 and has_previous:
                continue
            if export_stl:
                air_idx = list(device.allowed_permittivities).index(1)
                for idx in range(len(device.allowed_permittivities)):
                    if idx == air_idx and not export_air_stl:
                        continue
                    name = device.ordered_permittivity_tuples[idx][0]
                    export_stl_fn(
                        matrix=np.asarray(indices) == idx,
                        stl_filename=self.stl_dir / f"matrix_{iter_idx}_{device.name}_{name}.stl",
                        voxel_grid_size=device.single_voxel_grid_shape,
                    )
                if len(device.allowed_permittivities) > 2:
                    export_stl_fn(
                        matrix=np.asarray(indices) != air_idx,
                        stl_filename=self.stl_dir / f"matrix_{iter_idx}_{device.name}_non_air.stl",
                        voxel_grid_size=device.single_voxel_grid_shape,
                    )

            # image of indices
            if export_figure:
                fig = device_matrix_index_figure(
                    indices,
                    tuple(device.ordered_permittivity_tuples),
                )
                self.savefig(
                    self.cwd / "device",
                    f"matrix_indices_{iter_idx}_{device.name}.png",
                    fig,
                    dpi=72,
                )

            # raw permittivities and parameters
            jnp.save(self.params_dir / f"matrix_{iter_idx}_{device.name}.npy", indices)
            for k, v in device_params.items():
                jnp.save(self.params_dir / f"params_{iter_idx}_{device.name}_{k}.npy", v)

        return changed_voxels
