"""
vimshottari_visualizations.py
=============================

Generates the two figures that appear in the published version of:

  Oshop, R. (2026). 'A Bayesian derivation of the Vimśottarī Daśā years:
  Model comparison of Keplerian and palindromic accounts.' Research Note.

Figure 1: Nested Bayesian decomposition (palindromically-informed framing).
  Two panels:
    (a) Cumulative log Bayes factor across the four nested models.
    (b) Per-step Bayes factors on a log scale with Jeffreys categories.

Figure 2: Cross-scheme replication of the Keplerian 1/v law.
  Two panels:
    (a) Residuals of 1/v law for each of the five planets in each of
        five classical Indian daśā schemes.
    (b) Bayes factor for 1/v law per scheme, log scale.

Both figures are saved as 300-DPI JPEGs.

Dependencies: numpy, scipy, matplotlib.
  Import vimshottari_calculations for the numerical inputs.

Usage:
  python3 vimshottari_visualizations.py [output_dir]
"""
from __future__ import annotations

import os
import sys
from math import log, exp

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

# Import the companion calculations module.
import vimshottari_calculations as vc


# =============================================================================
# MATPLOTLIB DEFAULTS — paper style
# =============================================================================

plt.rcParams.update({
    'savefig.dpi': 300,
    'savefig.format': 'jpg',
    'savefig.bbox': 'tight',
    'savefig.facecolor': 'white',
    'font.family': 'serif',
    'font.size': 10,
    'axes.spines.top': False,
    'axes.spines.right': False,
    'axes.linewidth': 0.8,
})

# Colour palette
COL_BLUE       = '#4A6FA5'
COL_DARK_BLUE  = '#2E4865'
COL_LIGHT_BLUE = '#A8C0DA'
COL_PURPLE     = '#7B5C9E'
COL_CORAL      = '#C97157'
COL_GOLD       = '#C8A34B'
COL_GREEN      = '#6B8F6B'
COL_GRAY       = '#7A7A7A'


# =============================================================================
# FIGURE 1 — Nested Bayesian decomposition (palindromic framing)
# =============================================================================

def figure_1(output_dir: str) -> None:
    """Two-panel decomposition:
       (a) cumulative waterfall of log BF across four nested models
       (b) per-step Bayes factors on a log scale with Jeffreys categories.
    """
    # Pull numerical inputs from the calculations module so the figure
    # always reflects the current values.
    log_p_M0        = -log(vc.support_M0())
    log_p_Mweak     = -log(vc.support_M_weak())
    log_p_Mpalprior = -log(vc.support_M_palprior())

    obs = np.array([vc.OBSERVED_VIMSHOTTARI[k]
                    for k in vc.OBSERVED_VIMSHOTTARI])
    pred_dict = vc.m1_predictions()
    pred = np.array([pred_dict[k] for k in vc.OBSERVED_VIMSHOTTARI])
    log_p_M1 = vc.log_marginal_M1(obs, pred)

    # Three log Bayes factors across the nested hierarchy
    log_bf_structural = log_p_Mweak - log_p_M0               # 60-60 + Ketu=Mars
    log_bf_palindromic = log_p_Mpalprior - log_p_Mweak       # palindromic priors
    log_bf_keplerian  = log_p_M1 - log_p_Mpalprior           # 1/v law

    fig, axes = plt.subplots(1, 2, figsize=(14, 5.8),
                             gridspec_kw={'width_ratios': [1, 1.15],
                                          'wspace': 0.55})

    # -------------------------------------------------------------------------
    # Panel (a): Waterfall
    # -------------------------------------------------------------------------
    ax = axes[0]
    steps = [
        ('M₀\n(uniform null)',   0.00,                None),
        ('+60–60\n+Ketu=Mars',   log_bf_structural,   'structural'),
        ('+palindromic\npriors', log_bf_palindromic,  'palindromic'),
        ('+1/v\nKeplerian law',  log_bf_keplerian,    'physical'),
    ]
    cum = np.cumsum([s[1] for s in steps])

    bar_width = 0.6
    for i, (name, step_log_bf, kind) in enumerate(steps):
        if i == 0:
            ax.bar(i, 0.3, width=bar_width, bottom=-0.15,
                   color=COL_LIGHT_BLUE, edgecolor='#666', linewidth=0.8)
            ax.text(i, -0.8, 'log BF = 0', ha='center', fontsize=9, color='#444')
            continue

        bottom = cum[i-1]
        height = step_log_bf
        colour = {'structural': COL_BLUE,
                  'palindromic': COL_PURPLE,
                  'physical': COL_CORAL}[kind]
        ax.bar(i, height, width=bar_width, bottom=bottom,
               color=colour, edgecolor='white', linewidth=1.5)

        step_bf = exp(step_log_bf)
        mid = bottom + height / 2

        if step_log_bf < 3.0:
            # Small bars: label above with leader
            ax.annotate(
                f'Δ log BF\n= {step_log_bf:.2f}\n×{step_bf:.1f}',
                xy=(i, bottom + height),
                xytext=(i, bottom + height + 1.3),
                fontsize=8.2, color=colour, fontweight='bold',
                ha='center', va='bottom',
                arrowprops=dict(arrowstyle='-', color=colour, lw=0.6))
        else:
            label = (f'Δ log BF\n= {step_log_bf:.2f}\n×{step_bf:,.0f}'
                     if step_bf >= 10
                     else f'Δ log BF\n= {step_log_bf:.2f}\n×{step_bf:.2f}')
            ax.text(i, mid, label, ha='center', va='center',
                    fontsize=8.5, color='white', fontweight='bold')

        if i > 1:
            ax.plot([i - 1 + bar_width/2, i - bar_width/2], [bottom, bottom],
                    '--', color='#888', linewidth=0.8, alpha=0.6)

    # Cumulative totals above intermediate bars
    for i in range(1, len(steps) - 1):
        total_bf = exp(cum[i])
        label = (f'{total_bf:,.0f}' if total_bf < 1e5
                 else f'{total_bf:.2e}')
        ax.text(i, cum[i] + 0.4, f'cum. BF = {label}',
                ha='center', fontsize=8, color='#444', style='italic')

    # Final cumulative BF
    final_bf = exp(cum[-1])
    ax.text(len(steps) - 1 + 0.5, cum[-1],
            f'cum. BF\n= {final_bf:.2e}',
            ha='left', va='center', fontsize=8.5, color='#222',
            style='italic', fontweight='bold')

    # Axes
    ax.set_xticks(np.arange(len(steps)))
    ax.set_xticklabels([s[0] for s in steps], fontsize=9.5)
    ax.set_ylabel('Cumulative log Bayes factor (M vs M₀)', fontsize=10)
    ax.set_ylim(-1.5, 22)
    ax.set_xlim(-0.6, 3.7)
    ax.axhline(y=0, color='black', linewidth=0.6)

    # Jeffreys threshold markers
    for thresh, label in [(log(3),   'substantial'),
                          (log(10),  'strong'),
                          (log(30),  'very strong'),
                          (log(100), 'decisive')]:
        ax.axhline(y=thresh, color=COL_GRAY, linewidth=0.4,
                   linestyle=':', alpha=0.5)
        ax.text(3.55, thresh, label, ha='right', va='center',
                fontsize=7.5, color=COL_GRAY, style='italic')

    ax.set_title('(a) Cumulative evidence across nested models',
                 fontsize=10.5, pad=8)

    legend_items = [
        Patch(facecolor=COL_BLUE,   label='structural'),
        Patch(facecolor=COL_PURPLE, label='palindromic (Melakarta)'),
        Patch(facecolor=COL_CORAL,  label='physical (1/v)'),
    ]
    ax.legend(handles=legend_items, loc='upper left',
              frameon=False, fontsize=8.5)

    # -------------------------------------------------------------------------
    # Panel (b): Per-step Bayes factors, horizontal bar chart
    # -------------------------------------------------------------------------
    ax = axes[1]
    step_comparisons = [
        ('60–60 + Ketu=Mars\nvs uniform null',        exp(log_bf_structural),   'structural'),
        ('Palindromic priors\nvs structural',         exp(log_bf_palindromic),  'palindromic'),
        ('1/v Keplerian law\nvs structural+palindr.', exp(log_bf_keplerian),    'physical'),
    ]
    y_pos = np.arange(len(step_comparisons))

    for i, (name, bf, kind) in enumerate(step_comparisons):
        colour = {'structural': COL_BLUE,
                  'palindromic': COL_PURPLE,
                  'physical': COL_CORAL}[kind]
        log_bf = np.log10(bf)
        ax.barh(i, log_bf, color=colour, edgecolor='white', linewidth=1.5,
                height=0.55)

        if bf >= 1000:
            bf_str = f'{bf:,.0f}'
        elif bf >= 10:
            bf_str = f'{bf:.1f}'
        else:
            bf_str = f'{bf:.2f}'
        ax.text(log_bf + 0.12, i, f'BF = {bf_str}',
                va='center', fontsize=10, fontweight='bold')

    for x in [np.log10(3), np.log10(10), np.log10(30), np.log10(100)]:
        ax.axvline(x=x, color=COL_GRAY, linestyle=':',
                   linewidth=0.5, alpha=0.5)

    # Jeffreys category labels beneath axis, staggered
    ax.text(np.log10(1.5),  -0.85, 'barely worth\nmentioning',
            fontsize=7.5, color=COL_GRAY, style='italic',
            ha='center', va='top')
    ax.text(np.log10(5.5),  -1.15, 'substantial',
            fontsize=7.5, color=COL_GRAY, style='italic',
            ha='center', va='top')
    ax.text(np.log10(18),   -0.85, 'strong',
            fontsize=7.5, color=COL_GRAY, style='italic',
            ha='center', va='top')
    ax.text(np.log10(55),   -1.15, 'very strong',
            fontsize=7.5, color=COL_GRAY, style='italic',
            ha='center', va='top')
    ax.text(np.log10(1500), -0.85, 'decisive',
            fontsize=7.5, color=COL_GRAY, style='italic',
            ha='center', va='top')

    ax.set_yticks(y_pos)
    ax.set_yticklabels([s[0] for s in step_comparisons], fontsize=9.5)
    ax.invert_yaxis()

    ax.set_xlim(-0.3, 6.0)
    ax.set_ylim(-1.4, 2.8)
    ax.set_xticks([0, 1, 2, 3, 4, 5])
    ax.set_xticklabels(['1', '10', '100', '1,000', '10,000', '100,000'])
    ax.set_xlabel('Bayes factor (log scale)', fontsize=10, labelpad=28)

    ax.set_title('(b) Per-step Bayes factors and Jeffreys categories',
                 fontsize=10.5, pad=8)

    plt.suptitle('Figure 1.  Nested Bayesian decomposition '
                 '(palindromically-informed framing)',
                 fontsize=11.5, y=1.02)

    out_path = os.path.join(output_dir, 'figure_1_bayesian_decomposition.jpg')
    plt.savefig(out_path)
    plt.close()
    print(f"Figure 1 saved: {out_path}")


# =============================================================================
# FIGURE 2 — Cross-scheme replication
# =============================================================================

def figure_2(output_dir: str) -> None:
    """Two-panel cross-scheme figure:
       (a) residuals of 1/v law across five schemes
       (b) Bayes factor per scheme on a log scale.
    """
    results = [
        vc.analyse_scheme(name, s['total'], s['dashas'])
        for name, s in vc.COMPARISON_SCHEMES.items()
    ]

    # Reconstruct residuals for each scheme (in planet order)
    for r in results:
        scheme = vc.COMPARISON_SCHEMES[r['name']]
        preds = vc.one_over_v_predictions(r['inner_sum'], r['outer_sum'])
        planets = ['Mercury', 'Venus', 'Mars', 'Jupiter', 'Saturn']
        obs_arr = np.array([scheme['dashas'][p] for p in planets], dtype=float)
        pred_arr = np.array([preds[p] for p in planets])
        r['residuals'] = (pred_arr - obs_arr).tolist()

    fig, axes = plt.subplots(1, 2, figsize=(14, 5.5),
                             gridspec_kw={'width_ratios': [1.2, 1.0],
                                          'wspace': 0.35})

    # -------------------------------------------------------------------------
    # Panel (a): Residuals per planet per scheme, grouped bar chart
    # -------------------------------------------------------------------------
    ax = axes[0]
    planet_labels = ['Mer', 'Ven', 'Mar', 'Jup', 'Sat']
    n_schemes = len(results)
    bar_width = 0.13
    colours = [COL_DARK_BLUE, COL_BLUE, COL_GRAY, COL_GOLD, COL_GREEN]

    for i, r in enumerate(results):
        offset = (i - n_schemes/2 + 0.5) * bar_width
        x_positions = np.arange(5) + offset
        ax.bar(x_positions, r['residuals'], width=bar_width,
               color=colours[i], edgecolor='white', linewidth=0.6,
               label=r['name'])

    ax.axhline(y=0, color='black', linewidth=0.6)
    ax.set_xticks(np.arange(5))
    ax.set_xticklabels(planet_labels)
    ax.set_ylabel('Keplerian 1/v residual (years)')
    ax.set_xlabel('Planet')
    ax.set_title('(a) Residuals of 1/v law per scheme')
    ax.legend(loc='upper left', frameon=False, fontsize=8.5)
    ax.grid(True, axis='y', alpha=0.3, linewidth=0.4)

    # -------------------------------------------------------------------------
    # Panel (b): Per-scheme Bayes factor, horizontal log bars
    # -------------------------------------------------------------------------
    ax = axes[1]
    y_pos = np.arange(n_schemes)

    for i, r in enumerate(results):
        bf = r['BF']
        log_bf = np.log10(bf) if bf > 0 else -10
        if log_bf >= 0:
            ax.barh(i, log_bf, color=colours[i], edgecolor='white',
                    linewidth=1.0, height=0.55)
        else:
            ax.barh(i, log_bf, color=COL_CORAL, edgecolor='white',
                    linewidth=1.0, height=0.55, alpha=0.7)

        if bf >= 1000:
            bf_str = f'{bf:,.0f}'
        elif bf >= 10:
            bf_str = f'{bf:.1f}'
        else:
            bf_str = f'{bf:.2e}' if bf < 1 else f'{bf:.2f}'

        if log_bf >= 0:
            ax.text(log_bf + 0.08, i, f'BF = {bf_str}',
                    va='center', fontsize=9, fontweight='bold')
        else:
            ax.text(log_bf - 0.08, i, f'BF = {bf_str}',
                    va='center', fontsize=9, fontweight='bold', ha='right')

    for x in [np.log10(3), np.log10(10), np.log10(30), np.log10(100)]:
        ax.axvline(x=x, color=COL_GRAY, linestyle=':',
                   linewidth=0.5, alpha=0.5)
    ax.axvline(x=0, color='black', linewidth=0.8)

    ax.set_yticks(y_pos)
    ax.set_yticklabels([r['name'] for r in results], fontsize=9.5)
    ax.invert_yaxis()
    ax.set_xlabel('Bayes factor for 1/v law (log scale)', labelpad=6)
    ax.set_xlim(-10, 4)
    ax.set_ylim(n_schemes + 0.3, -0.5)  # room for annotation below bars
    ax.set_xticks([-9, -6, -3, 0, 1, 2, 3])
    ax.set_xticklabels(['10⁻⁹', '10⁻⁶', '10⁻³', '1', '10', '100', '1000'])

    # 'BF = 1' annotation placed below the bars so it doesn't collide with title
    ax.text(0, n_schemes - 0.6, 'BF = 1\n(null indifference)',
            fontsize=7.5, color='#444', style='italic',
            ha='center', va='top')

    ax.set_title('(b) Keplerian Bayes factor per scheme')

    plt.suptitle('Figure 2.  Cross-scheme replication of the Keplerian 1/v law',
                 fontsize=11.5, y=1.01)

    out_path = os.path.join(output_dir, 'figure_2_cross_scheme.jpg')
    plt.savefig(out_path)
    plt.close()
    print(f"Figure 2 saved: {out_path}")


# =============================================================================
# MAIN
# =============================================================================

def main() -> None:
    out_dir = sys.argv[1] if len(sys.argv) > 1 else '.'
    os.makedirs(out_dir, exist_ok=True)
    print(f"Saving figures to: {os.path.abspath(out_dir)}")
    print("-" * 60)
    figure_1(out_dir)
    figure_2(out_dir)
    print("-" * 60)
    print("Both figures complete.")


if __name__ == '__main__':
    main()