#!/usr/bin/env python3
"""pyGDIS 1.0.0: reproducible logistic-map tutorial with rich figures.

Run: python pygdis_logistic_tutorial.py
Output: tutorial_output/*.csv, *.png, *.pdf
"""
from pathlib import Path
import argparse
import shutil
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from gdis import GDIS

OUT = Path(__file__).resolve().parent / 'tutorial_output'
OUT.mkdir(exist_ok=True)
R = np.linspace(2.8, 4.0, 121)
N_TRANSIENT = 600
N_SAMPLES = 500
X0 = 0.314159


def logistic_trajectory(r, x0=X0, n_transient=N_TRANSIENT, n_samples=N_SAMPLES):
    """Simulate x[n+1] = r*x[n]*(1-x[n]); discard the transient."""
    x = float(x0)
    for _ in range(n_transient):
        x = r * x * (1.0 - x)
    out = np.empty(n_samples, dtype=float)
    for k in range(n_samples):
        x = r * x * (1.0 - x)
        out[k] = x
    return out


def plot_bifurcation(trajectories, parameters):
    fig, ax = plt.subplots(figsize=(11, 5.5), layout='constrained')
    for r, t in zip(parameters, trajectories):
        ax.scatter(np.full(125, r), t[-125:, 0], s=0.25, c='#2563eb', alpha=0.5, rasterized=True)
    ax.set(xlabel='Control parameter r', ylabel='State x', title='Logistic map | Bifurcation diagram')
    ax.axvline(3.56995, c='#f97316', ls='--', lw=1.8, label='Chaos accumulation point ≈ 3.56995')
    ax.legend(frameon=False, fontsize=9)
    fig.savefig(OUT / '01_bifurcation.png', dpi=220)
    fig.savefig(OUT / '01_bifurcation.pdf')
    plt.close(fig)


def plot_trajectories(trajectories, parameters):
    targets = [2.9, 3.2, 3.5, 3.9]
    fig, axes = plt.subplots(2, 2, figsize=(11, 6.7), sharex=True, sharey=True, layout='constrained')
    for ax, target in zip(axes.flat, targets):
        idx = int(np.argmin(abs(parameters - target)))
        ax.plot(np.arange(90), trajectories[idx][:90, 0], lw=1.6, c='#2563eb')
        ax.set_title(f'r = {parameters[idx]:.2f}', loc='left', weight='bold')
        ax.grid(alpha=0.17)
        ax.set_ylim(-.03, 1.03)
    for ax in axes[-1]: ax.set_xlabel('Iteration (post-transient)')
    for ax in axes[:, 0]: ax.set_ylabel('x')
    fig.suptitle('Stable, periodic and chaotic-looking trajectories', fontsize=15, weight='bold')
    fig.savefig(OUT / '02_trajectories.png', dpi=220)
    fig.savefig(OUT / '02_trajectories.pdf')
    plt.close(fig)


def plot_gdis(result):
    fig, ax = plt.subplots(figsize=(11, 5.5), layout='constrained')
    ax.plot(result.parameters, result.gdis, c='#2563eb', lw=2.6, label='GDIS')
    ax.fill_between(result.parameters, result.gdis, color='#93c5fd', alpha=0.32)
    ax.plot(result.parameters, result.sustained_instability, c='#0d9488', lw=1.8, label='Sustained instability')
    ax.axvline(3.56995, c='#f97316', ls='--', lw=1.7, label='Known logistic-map reference ≈3.56995')
    ax.set(xlabel='Control parameter r', ylabel='Bounded score', title='Generalized Dynamical Instability Score')
    ax.set_ylim(0, 1.02)
    ax.grid(alpha=0.18)
    ax.legend(frameon=False, fontsize=9)
    fig.savefig(OUT / '03_gdis_profile.png', dpi=220)
    fig.savefig(OUT / '03_gdis_profile.pdf')
    plt.close(fig)


def plot_components(result):
    c = result.components
    fig, axes = plt.subplots(2, 2, figsize=(12, 8), layout='constrained')
    ax = axes[0, 0]
    for name, color in [('jacobian_saturated', '#2563eb'), ('stretching_saturated', '#0d9488'), ('expansion_saturated', '#9333ea')]:
        ax.plot(result.parameters, c[name], lw=1.8, label=name.replace('_saturated', '').title(), color=color)
    ax.set_title('Normalized core descriptors')
    ax.legend(frameon=False)
    ax = axes[0, 1]
    ax.plot(result.parameters, c['entropy_raw'], lw=2, color='#d97706', label='H: nonlinear complexity')
    ax.plot(result.parameters, c['temporal_raw'], lw=2, color='#7c3aed', label='T: temporal persistence')
    ax.set_title('Additional descriptors (raw)')
    ax.legend(frameon=False)
    ax = axes[1, 0]
    ax.plot(result.parameters, c['transition_energy'], lw=2.1, color='#e11d48', label='Transition energy')
    ax.plot(result.parameters, c['critical_window'], lw=1.6, color='#64748b', label='Localization window')
    ax.set_title('Localized transition evidence')
    ax.legend(frameon=False)
    ax = axes[1, 1]
    ax.plot(result.parameters, result.gdis, lw=2.3, color='#2563eb', label='Final GDIS')
    ax.plot(result.parameters, result.sustained_instability, lw=1.9, color='#0d9488', label='Sustained')
    ax.plot(result.parameters, result.transition_instability, lw=1.8, color='#e11d48', label='Weighted transition term')
    ax.set_title('Final score and constituents')
    ax.legend(frameon=False)
    for ax in axes.flat:
        ax.set_xlabel('r')
        ax.grid(alpha=0.18)
    fig.suptitle('pyGDIS | Descriptor and transition diagnostics', fontsize=16, weight='bold')
    fig.savefig(OUT / '04_components.png', dpi=220)
    fig.savefig(OUT / '04_components.pdf')
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description="Generate pyGDIS logistic-map figures")
    parser.add_argument("--publish-static-dir", type=Path, default=None,
                        help="Optional Django static destination for the generated PNG figures")
    args = parser.parse_args()
    plt.rcParams.update({'font.size': 11, 'axes.titlesize': 13, 'axes.spines.top': False,
                         'axes.spines.right': False, 'savefig.facecolor': 'white'})
    # Data-only mode: each (N,1) array is a separate trajectory for a unique r.
    trajectories = [logistic_trajectory(r)[:, None] for r in R]
    long_df = pd.DataFrame({
        'parameter': np.repeat(R, N_SAMPLES),
        'time': np.tile(np.arange(N_SAMPLES), len(R)),
        'x1': np.concatenate([t[:, 0] for t in trajectories])
    })
    long_df.to_csv(OUT / 'synthetic_logistic_data.csv', index=False)

    # Do not provide a critical_value: allow pyGDIS to select one from transition energy.
    result = GDIS().fit_transform(trajectories=trajectories, parameters=R)
    df = result.to_dataframe()
    df.to_csv(OUT / 'gdis_results.csv', index=False)

    plot_bifurcation(trajectories, R)
    plot_trajectories(trajectories, R)
    plot_gdis(result)
    plot_components(result)

    peak = int(np.argmax(result.gdis))
    summary = (
        'pyGDIS 1.0.0 | Logistic map example\n'
        f'Parameter sweep: {len(R)} parameter settings from {R.min():.2f} to {R.max():.2f}\n'
        f'Trajectory length: {N_SAMPLES}, transient discarded: {N_TRANSIENT}\n'
        f'Maximum GDIS: {result.gdis[peak]:.5f} at r = {R[peak]:.4f}\n'
        f"Transition center: {result.metadata['critical_value']:.5f} "
        f"({result.metadata['critical_value_source']})\n"
        'Important: GDIS is a descriptive score, not a proof of chaos or an exact bifurcation detector.\n'
    )
    (OUT / 'analysis_summary.txt').write_text(summary, encoding='utf-8')
    print(summary)
    print(f'CSV data, results, and four figures (PNG+PDF) saved to {OUT}')
    if args.publish_static_dir is not None:
        args.publish_static_dir.mkdir(parents=True, exist_ok=True)
        for png in OUT.glob('*.png'):
            shutil.copy2(png, args.publish_static_dir / png.name)
        print(f'Published PNG figures to: {args.publish_static_dir}')


if __name__ == '__main__':
    main()