#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Plot the data generated by nonunitary_magnetization_calculation.py."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import json
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
from matplotlib.ticker import FuncFormatter


@dataclass(frozen=True)
class PlotParams:
    data_root: str = "nonunitary_magnetization_results/data"
    figure_root: str = "nonunitary_magnetization_results/figures"
    global_font_size: float = 14.0
    legend_font_size: float = 11.5
    line_width: float = 2.2
    marker_size: float = 5.5
    dpi: int = 300
    zero_threshold: float = 1.0e-10


COLORS = {
    1.0: "#4C72B0",
    0.6: "#55A868",
    0.25: "#C44E52",
    0.0: "#8172B2",
}


def configure_plotting(p: PlotParams) -> None:
    plt.rcParams.update(
        {
            "font.family": "serif",
            "font.serif": ["Times New Roman", "Times", "STIXGeneral", "DejaVu Serif"],
            "mathtext.fontset": "stix",
            "font.size": p.global_font_size,
            "axes.labelsize": p.global_font_size,
            "xtick.labelsize": p.global_font_size,
            "ytick.labelsize": p.global_font_size,
            "legend.fontsize": p.legend_font_size,
            "axes.linewidth": 1.0,
            "xtick.direction": "in",
            "ytick.direction": "in",
            "xtick.top": True,
            "ytick.right": True,
            "savefig.dpi": p.dpi,
            "figure.dpi": p.dpi,
        }
    )


def compact_number(x: float, pos: int | None = None) -> str:
    if abs(x) < 5.0e-10:
        return "0"
    return f"{x:.3g}"


def save_line_plot(
    x_series: list[np.ndarray],
    y_series: list[np.ndarray],
    labels: list[str],
    colors: list[str],
    linestyles: list[str],
    markers: list[str],
    xlabel: str,
    ylabel: str,
    output: Path,
    p: PlotParams,
    y_min_zero_extension: bool = False,
) -> None:
    fig, ax = plt.subplots(figsize=(5.2, 5.2))
    ax.set_box_aspect(1)

    for x, y, label, color, ls, marker in zip(
        x_series, y_series, labels, colors, linestyles, markers
    ):
        stride = max(1, len(x) // 12)
        ax.plot(
            x,
            y,
            label=label,
            color=color,
            linestyle=ls,
            linewidth=p.line_width,
            marker=marker,
            markersize=p.marker_size,
            markevery=stride,
        )

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.xaxis.set_major_formatter(FuncFormatter(compact_number))
    ax.yaxis.set_major_formatter(FuncFormatter(compact_number))

    if y_min_zero_extension:
        ymax = max(float(np.max(y)) for y in y_series)
        pad = max(0.04 * ymax, 1.0e-5)
        ax.set_ylim(-pad, 1.08 * ymax if ymax > 0 else pad)

    ax.legend(frameon=False)
    fig.tight_layout()
    fig.savefig(output, bbox_inches="tight")
    plt.close(fig)


def plot_all() -> None:
    p = PlotParams()
    configure_plotting(p)

    data_dir = Path(p.data_root)
    fig_dir = Path(p.figure_root)
    fig_dir.mkdir(parents=True, exist_ok=True)

    scan = np.load(data_dir / "temperature_scan.npz")
    data = scan["data"]
    columns = [str(x) for x in scan["columns"]]
    idx = {name: i for i, name in enumerate(columns)}
    r_values = scan["r_values"]

    x_series = []
    m_series = []
    q_series = []
    labels = []
    colors = []
    linestyles = []
    markers = ["o", "s", "^", "D"]

    for j, r in enumerate(r_values):
        subset = data[np.isclose(data[:, idx["r"]], r)]
        x_series.append(subset[:, idx["T_over_Tc"]])
        m_series.append(subset[:, idx["M_z_muB_per_site"]])
        q_series.append(subset[:, idx["q_gap_average"]])
        labels.append(rf"$r={r:g}$")
        colors.append(COLORS.get(float(r), None))
        linestyles.append("-" if r < 1.0 else "--")

    save_line_plot(
        x_series,
        m_series,
        labels,
        colors,
        linestyles,
        markers,
        r"$T/T_c$",
        r"$M_z\;(\mu_B/{\rm site})$",
        fig_dir / "magnetization_vs_temperature.png",
        p,
        y_min_zero_extension=True,
    )

    save_line_plot(
        x_series,
        q_series,
        labels,
        colors,
        linestyles,
        markers,
        r"$T/T_c$",
        r"$\langle q_{\Delta,z}\rangle$",
        fig_dir / "pair_spin_vs_temperature.png",
        p,
        y_min_zero_extension=True,
    )

    r_scan = np.loadtxt(data_dir / "nonunitarity_scan.dat")
    fig, ax = plt.subplots(figsize=(5.2, 5.2))
    ax.set_box_aspect(1)
    ax.plot(
        r_scan[:, 1],
        r_scan[:, 3],
        color="#C44E52",
        linewidth=p.line_width,
        marker="o",
        markersize=p.marker_size,
        markevery=8,
        label=r"$M_z$",
    )
    ax.set_xlabel(r"$\eta=1-r^2$")
    ax.set_ylabel(r"$M_z\;(\mu_B/{\rm site})$")
    ax.xaxis.set_major_formatter(FuncFormatter(compact_number))
    ax.yaxis.set_major_formatter(FuncFormatter(compact_number))
    ymax = float(np.max(r_scan[:, 3]))
    ax.set_ylim(-0.04 * ymax, 1.08 * ymax)
    ax.legend(frameon=False)
    fig.tight_layout()
    fig.savefig(
        fig_dir / "magnetization_vs_nonunitarity.png",
        bbox_inches="tight",
    )
    plt.close(fig)

    maps = np.load(data_dir / "representative_maps.npz")
    k = maps["k"] / np.pi
    xi = maps["xi"]
    r0 = float(maps["r"])
    T0 = float(maps["T"])
    extent = [-1.0, 1.0, -1.0, 1.0]

    for key, label, filename in [
        (
            "q_gap_map",
            r"$q_{\Delta,z}(\boldsymbol{k})$",
            "q_gap_map.png",
        ),
        (
            "delta_n_map",
            r"$n_{\uparrow}(\boldsymbol{k})-n_{\downarrow}(\boldsymbol{k})$",
            "spin_density_difference_map.png",
        ),
    ]:
        z = maps[key]
        vmax = float(np.max(np.abs(z)))
        fig, ax = plt.subplots(figsize=(5.1, 5.1))
        ax.set_box_aspect(1)
        im = ax.imshow(
            z.T,
            origin="lower",
            extent=extent,
            interpolation="spline16",
            cmap="coolwarm",
            norm=TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax),
            aspect="equal",
        )
        ax.contour(
            k,
            k,
            xi.T,
            levels=[0.0],
            colors="#39FF14",
            linewidths=1.4,
        )
        ax.set_xlabel(r"$k_x/\pi$")
        ax.set_ylabel(r"$k_y/\pi$")
        ax.set_xticks([-1, 0, 1])
        ax.set_yticks([-1, 0, 1])
        cbar = fig.colorbar(
            im,
            ax=ax,
            pad=0.025,
            fraction=0.042,
            shrink=0.86,
            aspect=25,
            extend="both",
        )
        cbar.set_label(label)
        cbar.ax.yaxis.set_major_formatter(FuncFormatter(compact_number))
        fig.tight_layout()
        fig.savefig(fig_dir / filename, bbox_inches="tight")
        plt.close(fig)

    dos = np.loadtxt(data_dir / "representative_spin_dos.dat")
    fig, ax = plt.subplots(figsize=(5.2, 5.2))
    ax.set_box_aspect(1)
    ax.plot(
        dos[:, 0],
        dos[:, 1],
        color="#D9534F",
        linewidth=p.line_width,
        label=r"$N_{\uparrow}(E)$",
    )
    ax.plot(
        dos[:, 0],
        dos[:, 2],
        color="#3B6FB6",
        linewidth=p.line_width,
        linestyle="--",
        label=r"$N_{\downarrow}(E)$",
    )
    ax.set_xlabel(r"$E/t$")
    ax.set_ylabel(r"${\rm DOS}$")
    ax.xaxis.set_major_formatter(FuncFormatter(compact_number))
    ax.yaxis.set_major_formatter(FuncFormatter(compact_number))
    ax.legend(frameon=False)
    fig.tight_layout()
    fig.savefig(fig_dir / "spin_resolved_dos.png", bbox_inches="tight")
    plt.close(fig)

    print(
        f"Figures saved in {fig_dir}. Representative state: "
        f"r={r0:g}, T={T0:.6f}."
    )


if __name__ == "__main__":
    plot_all()
