"""Visualization utilities for NDVI curves and prediction GIFs."""

from __future__ import annotations

import math
import os
import shutil

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

from ndvi_forecast._utils import to_numpy, smooth1d


def plot_ndvi_curves(
    ndvi_array: np.ndarray,
    year_label: str,
    dap_index: int = 6,
    group_window: int = 3,
    save_dir: str | None = None,
) -> plt.Figure:
    """Plot all NDVI curves with a group-stats DAP window highlighted."""
    seq_len = ndvi_array.shape[1]
    x = np.arange(1, seq_len + 1)

    fig, ax = plt.subplots(figsize=(12, 6))
    for curve in ndvi_array:
        ax.plot(x, curve, alpha=0.6)

    ax.set_title(f"NDVI Curves \u2013 {year_label}")
    ax.set_xlabel("Day After Planting (DAP)")
    ax.set_ylabel("NDVI")
    ax.set_ylim(0, 1.05)
    ax.axvspan(dap_index + 1, dap_index + 1 + group_window, color="red", alpha=0.2, label="Grouping Window")
    ax.grid(True)
    fig.tight_layout()

    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        fig.savefig(os.path.join(save_dir, f"ndvi_{year_label}.svg"))

    return fig


def plot_avg_ndvi(
    ndvi_array: np.ndarray,
    year_label: str,
    dap_index: int = 6,
    group_window: int = 3,
    save_dir: str | None = None,
) -> plt.Figure:
    """Plot average NDVI curve with group-stats window."""
    avg = ndvi_array.mean(axis=0)
    seq_len = avg.shape[0]
    x = np.arange(1, seq_len + 1)

    fig, ax = plt.subplots(figsize=(12, 6))
    ax.plot(x, avg, alpha=0.6)
    ax.set_title(f"Average NDVI \u2013 {year_label}")
    ax.set_xlabel("Day After Planting (DAP)")
    ax.set_ylabel("NDVI")
    ax.set_ylim(0, 1.05)
    ax.axvspan(dap_index + 1, dap_index + 1 + group_window, color="red", alpha=0.2, label="Grouping Window")
    ax.grid(True)
    fig.tight_layout()

    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        fig.savefig(os.path.join(save_dir, f"ndvi_avg_{year_label}.svg"))

    return fig


def plot_overlapped_ndvi(
    ndvi_dict: dict[str, np.ndarray],
    dap_index: int = 6,
    group_window: int = 3,
    save_dir: str | None = None,
) -> plt.Figure:
    """Overlay average NDVI curves from multiple years/sites.

    Args:
        ndvi_dict: mapping of label -> (N, T) NDVI array.
    """
    fig, ax = plt.subplots(figsize=(12, 6))

    for label, arr in ndvi_dict.items():
        avg = arr.mean(axis=0)
        x = np.arange(1, avg.shape[0] + 1)
        ax.plot(x, avg, label=label, alpha=0.7)

    ax.axvspan(dap_index + 1, dap_index + 1 + group_window, color="red", alpha=0.2, label="Grouping Window")
    ax.set_title("NDVI Curve Overlap")
    ax.set_xlabel("Day After Planting (DAP)")
    ax.set_ylabel("NDVI")
    ax.set_ylim(0, 1.05)
    ax.grid(True)
    ax.legend()
    fig.tight_layout()

    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        fig.savefig(os.path.join(save_dir, "ndvi_overlapped.svg"))

    return fig


def create_prediction_gif(
    pred_chunks: np.ndarray,
    full_ndvi: np.ndarray,
    input_len: int,
    filename: str = "ndvi_forecast.gif",
    base_dir: str = "./results/",
    n_curves: int = 10,
    step_duration: int = 1000,
    smooth_preds: bool = True,
    smooth_kind: str = "ma",
    smooth_w: int = 10,
    dpi: int = 600,
    svg_snapshots_daps: tuple[int, ...] = (65, 80, 100),
    svg_dir: str = "./results/snapshots",
) -> str:
    """Create an animated GIF of autoregressive predictions.

    Returns the path to the saved GIF.
    """
    pred_chunks = np.asarray(pred_chunks)
    if pred_chunks.ndim == 4:
        pred_chunks = pred_chunks.squeeze(-1)
    full_ndvi = np.asarray(full_ndvi)

    os.makedirs(base_dir, exist_ok=True)
    os.makedirs(svg_dir, exist_ok=True)

    save_path = os.path.join(base_dir, filename)
    N, num_chunks, step = pred_chunks.shape
    total_len = full_ndvi.shape[1]
    frames = []
    pending_daps = set(int(d) for d in svg_snapshots_daps if 0 <= d < total_len)

    for chunk_idx in range(num_chunks):
        fig, ax = plt.subplots(figsize=(7, 4))
        latest_dap = input_len - 1

        for i in range(min(N, n_curves)):
            full_gt = to_numpy(full_ndvi[i])
            input_seq = full_gt[:input_len]
            future_gt = full_gt[input_len:]

            pred_so_far = pred_chunks[i, : chunk_idx + 1].reshape(-1)
            if smooth_preds and pred_so_far.size > 0:
                pred_so_far = smooth1d(pred_so_far, kind=smooth_kind, w=smooth_w)

            ax.plot(range(input_len), input_seq, color="blue", alpha=1.0, label="Input" if i == 0 else "")
            ax.plot(range(input_len, total_len), future_gt, color="green", alpha=1.0, label="GT" if i == 0 else "")

            if pred_so_far.size > 0:
                x_all = np.arange(input_len, input_len + len(pred_so_far))
                ax.plot(x_all, pred_so_far, color="red", alpha=1.0, zorder=2, label="Pred" if i == 0 else "")
                take = min(step, len(pred_so_far))
                if take > 0:
                    ax.plot(x_all[-take:], pred_so_far[-take:], color="darkred", alpha=1.0, zorder=3)
                latest_dap = max(latest_dap, input_len + len(pred_so_far) - 1)

        ax.set_xlim(0, total_len)
        ax.set_ylim(0, 1.05)
        ax.set_title(f"Autoregressive Prediction \u2013 Frame {chunk_idx + 1}/{num_chunks}")
        ax.grid(True, alpha=0.3)
        fig.tight_layout()

        frame_path = f"/tmp/frame_{chunk_idx}.png"
        fig.savefig(frame_path, dpi=dpi, bbox_inches="tight", facecolor="white")
        frames.append(Image.open(frame_path).convert("P"))

        for d in sorted(d for d in pending_daps if d <= latest_dap):
            svg_path = os.path.join(svg_dir, f"snapshot_dap{d:03d}.svg")
            fig.savefig(svg_path, format="svg", bbox_inches="tight")
            pending_daps.discard(d)

        plt.close(fig)

    if frames:
        frames[0].save(save_path, save_all=True, append_images=frames[1:],
                       duration=step_duration, loop=0, optimize=True, disposal=2)

    return save_path
