#!/usr/bin/env python3.10

"""Fit polynomial models to the CDI and Selic monthly series."""

from pathlib import Path
import shutil
import sys

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np


DEFAULT_SAMPLE_POINTS = np.arange(1.0, 71.0, dtype=float)
DEFAULT_CDI_VALUES = np.array(
    [
        0.6582, 0.5925, 0.7569, 0.6640, 0.7501, 0.7908, 0.8592, 0.8863, 0.8446, 0.8057,
        0.8056, 0.9272, 0.8606, 0.8424, 0.9189, 0.8388, 0.9853, 0.9527, 0.9666, 1.0724,
        0.9398, 0.8807, 0.8586, 0.9047, 0.8853, 0.7416, 0.8083, 0.6999, 0.7325, 0.6386,
        0.6754, 0.6866, 0.5372, 0.6073, 0.5445, 0.5342, 0.5866, 0.4816, 0.5377, 0.6008,
        0.5849, 0.5919, 0.7088, 0.6958, 0.6991, 0.8034, 0.7105, 0.7804, 0.8398, 0.7827,
        0.7600, 0.8154, 0.8583, 0.8174, 0.9404, 0.8595, 0.9006, 0.9448, 0.8379, 0.9558,
        0.9293, 0.8185, 1.0361, 0.9483, 0.9838, 1.0658, 1.1773, 1.1075, 1.1075, 1.1077,
    ],
    dtype=float,
)
DEFAULT_SELIC_VALUES = np.array(
    [
        0.6706, 0.6034, 0.7716, 0.6761, 0.7622, 0.8033, 0.8717, 0.8990, 0.8579, 0.8169,
        0.8169, 0.9400, 0.8724, 0.8531, 0.9302, 0.8490, 0.9983, 0.9661, 0.9778, 1.0851,
        0.9527, 0.8927, 0.8711, 0.9189, 0.9026, 0.7587, 0.8325, 0.7221, 0.7564, 0.6522,
        0.6917, 0.7040, 0.5489, 0.6228, 0.5593, 0.5593, 0.6105, 0.4992, 0.5548, 0.6188,
        0.6039, 0.6102, 0.7297, 0.7156, 0.7179, 0.8157, 0.7236, 0.7942, 0.8539, 0.7942,
        0.7697, 0.8266, 0.8700, 0.8284, 0.9532, 0.8700, 0.9116, 0.9550, 0.8463, 0.9653,
        0.9426, 0.8288, 1.0478, 0.9587, 0.9924, 1.0741, 1.1863, 1.1163, 1.1163, 1.1163,
    ],
    dtype=float,
)


def load_series(data_path):
    """Load one two-column time series file."""
    data = np.loadtxt(data_path)
    return data[:, 0], data[:, 1]


def load_input_series(command_line_arguments):
    """Load external data or fall back to the embedded default series."""
    if len(command_line_arguments) == 0:
        return DEFAULT_SAMPLE_POINTS.copy(), DEFAULT_CDI_VALUES.copy(), DEFAULT_SELIC_VALUES.copy()
    if len(command_line_arguments) != 2:
        raise ValueError("Pass either no input files or exactly two files: CDI and Selic.")

    cdi_path = Path(command_line_arguments[0]).expanduser().resolve()
    selic_path = Path(command_line_arguments[1]).expanduser().resolve()
    sample_points, cdi_values = load_series(cdi_path)
    sample_points_selic, selic_values = load_series(selic_path)

    if not np.allclose(sample_points, sample_points_selic):
        raise ValueError("The CDI and Selic input files must use the same sample points.")

    return sample_points, cdi_values, selic_values


def configure_plot_style():
    """Configure a consistent LaTeX-like plotting style."""
    mpl.rcParams.update(
        {
            "text.usetex": shutil.which("latex") is not None,
            "font.family": "serif",
            "font.serif": ["Computer Modern Roman", "DejaVu Serif"],
            "mathtext.fontset": "cm",
            "axes.titlesize": 15,
            "axes.labelsize": 13,
            "xtick.labelsize": 11,
            "ytick.labelsize": 11,
            "legend.fontsize": 11,
        }
    )


def main():
    """Fit degrees 1 through 7 and save one PDF for each degree."""
    configure_plot_style()

    sample_points, cdi_values, selic_values = load_input_series(sys.argv[1:])

    program_directory = Path(__file__).resolve().parent
    output_directory = program_directory.parent / "Figuras"
    output_directory.mkdir(exist_ok=True)

    evaluation_point = float(sample_points[-1] + 1.0)
    dense_grid = np.linspace(sample_points[0], evaluation_point, 400)

    for degree in range(1, 8):
        cdi_coefficients = np.polyfit(sample_points, cdi_values, degree)
        selic_coefficients = np.polyfit(sample_points, selic_values, degree)
        cdi_estimate = float(np.polyval(cdi_coefficients, evaluation_point))
        selic_estimate = float(np.polyval(selic_coefficients, evaluation_point))

        plt.figure(figsize=(7.0, 5.0))
        plt.plot(sample_points, cdi_values, "o", label="CDI data")
        plt.plot(sample_points, selic_values, "s", label="Selic data")
        plt.plot(dense_grid, np.polyval(cdi_coefficients, dense_grid), label=f"CDI degree {degree}")
        plt.plot(dense_grid, np.polyval(selic_coefficients, dense_grid), label=f"Selic degree {degree}")
        plt.axvline(evaluation_point, color="0.4", linewidth=1.0, linestyle="--")
        plt.xlabel("Month index")
        plt.ylabel("Rate")
        plt.grid(True, linewidth=0.4, alpha=0.5)
        plt.legend()
        plt.tight_layout()
        plt.savefig(output_directory / f"p6_pg{degree}.pdf", bbox_inches="tight")
        plt.close()

        print(f"degree = {degree}: CDI = {cdi_estimate:.10f}, Selic = {selic_estimate:.10f}")


if __name__ == "__main__":
    main()