from pathlib import Path
import os
import sys


SCRIPT_DIR = Path(__file__).resolve().parent
ROOT = SCRIPT_DIR.parent
MPLCONFIGDIR = ROOT / ".mplconfig_pic_fig03_simple"
MPLCONFIGDIR.mkdir(exist_ok=True)
for lock_file in MPLCONFIGDIR.glob("*.matplotlib-lock"):
    try:
        lock_file.unlink()
    except OSError:
        pass
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))
sys.path.insert(0, str(ROOT))

import matplotlib.pyplot as plt
import numpy as np

from lda.utils import gaspari_cohn


OUTDIR = SCRIPT_DIR
PNG_PATH = OUTDIR / "Fig03_gc_localization_physical_vs_spatial_latent.png"
PDF_PATH = OUTDIR / "Fig03_gc_localization_physical_vs_spatial_latent.pdf"


LATENT_SPACING = 16  # 2048 physical grid points / 128 spatial latent cells
RADII = [16, 32, 48, 64, 96, 128, 192, 256]


def main() -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Serif",
            "font.size": 11,
            "mathtext.fontset": "dejavuserif",
            "axes.linewidth": 0.9,
        }
    )

    fig, axes = plt.subplots(1, 2, figsize=(12.8, 5.2), facecolor="white", sharey=True)
    ax_phys, ax_lat = axes
    distances = np.linspace(0, max(RADII) * 1.08, 900)
    colors = plt.cm.viridis(np.linspace(0.06, 0.96, len(RADII)))

    for radius, color in zip(RADII, colors):
        weights = gaspari_cohn(distances, float(radius))
        ax_phys.plot(distances, weights, color=color, lw=2.4, label=f"R={radius}")

        latent_offsets = np.arange(0, int(np.ceil(radius / LATENT_SPACING)) + 2) * LATENT_SPACING
        latent_offsets = latent_offsets[latent_offsets <= max(RADII) * 1.08]
        latent_weights = gaspari_cohn(latent_offsets, float(radius))
        ax_lat.plot(
            latent_offsets / LATENT_SPACING,
            latent_weights,
            linestyle="-",
            linewidth=2.0,
            alpha=0.65,
            color=color,
        )
        ax_lat.plot(
            latent_offsets / LATENT_SPACING,
            latent_weights,
            linestyle="",
            marker="o",
            markersize=4.6,
            markerfacecolor=color,
            markeredgecolor=color,
            markeredgewidth=1.4,
        )

    ax_phys.set_xlim(0, max(RADII) * 1.06)
    ax_lat.set_xlim(0, max(RADII) / LATENT_SPACING * 1.06)
    ax_phys.set_ylim(-0.035, 1.10)
    ax_phys.set_xlabel("Distance from observation (physical grid points)")
    ax_lat.set_xlabel("Distance from observation (spatial latent cells)")
    ax_phys.set_ylabel("Localization weight")
    for ax in axes:
        ax.grid(True, color="0.88", lw=0.8)
    ax_phys.text(0.02, 0.96, "(a)", transform=ax_phys.transAxes, fontsize=14, ha="left", va="top")
    ax_lat.text(0.02, 0.96, "(b)", transform=ax_lat.transAxes, fontsize=14, ha="left", va="top")
    ax_phys.legend(
        title="Localization radius",
        ncol=2,
        frameon=False,
        loc="upper right",
        columnspacing=1.4,
        handlelength=2.2,
    )

    fig.tight_layout(w_pad=2.0)
    fig.savefig(PNG_PATH, dpi=300, bbox_inches="tight")
    fig.savefig(PDF_PATH, bbox_inches="tight")
    print(f"saved {PNG_PATH}")
    print(f"saved {PDF_PATH}")


if __name__ == "__main__":
    main()
