# -*- coding: utf-8 -*-
"""HWE Cognitive Warfare Simulation — Lux Ferox v3.0
Reflexive Index dynamic model with network topology and doctrinal oscillation."""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import networkx as nx
import pandas as pd
import os

matplotlib.rcParams['figure.figsize'] = (11, 6)
matplotlib.rcParams['font.size'] = 11
matplotlib.rcParams['axes.spines.top'] = False
matplotlib.rcParams['axes.spines.right'] = False
os.makedirs('figures', exist_ok=True)

# ── Colour palette ────────────────────────────────────────────────────────────
LUX_BLUE  = '#19387a'
LUX_GOLD  = '#b48c28'
LUX_RED   = '#8b1a1a'
LUX_GREEN = '#1a5c2e'
LUX_GRAY  = '#555555'

# ══════════════════════════════════════════════════════════════════════════════
# Core RI function — corrected v2.0 formulation
# ══════════════════════════════════════════════════════════════════════════════
def reflex_index(V, tau, C, F):
    """RI(t) = V(t)*tau(t) / (C(t)*F(t))  — with epsilon to avoid div/0."""
    return (V * tau) / (C * F + 1e-9)

# ══════════════════════════════════════════════════════════════════════════════
# FIGURE 1 — Baseline simulation
# ══════════════════════════════════════════════════════════════════════════════
np.random.seed(42)
T = 150
t = np.arange(T)

V   = 1.0 + 0.20 * t + 0.5  * np.random.randn(T).cumsum() * 0.1
tau = 0.3 + 0.005 * t + 0.02 * np.random.randn(T)
C   = 2.0 + 0.10 * t + 0.3  * np.random.randn(T).cumsum() * 0.05
F   = 1.5 + 0.15 * t + 0.3  * np.random.randn(T).cumsum() * 0.05

V   = np.maximum(V, 0.1)
tau = np.clip(tau, 0.01, 1.0)
C   = np.maximum(C, 0.1)
F   = np.maximum(F, 0.1)

RI = reflex_index(V, tau, C, F)

V   = 1.0 + 0.20 * t + 0.5  * np.random.randn(T).cumsum() * 0.1
tau = 0.3 + 0.005 * t + 0.02 * np.random.randn(T)
C   = 2.0 + 0.10 * t + 0.3  * np.random.randn(T).cumsum() * 0.05
F   = 1.5 + 0.15 * t + 0.3  * np.random.randn(T).cumsum() * 0.05

V   = np.maximum(V, 0.1)
tau = np.clip(tau, 0.01, 1.0)
C   = np.maximum(C, 0.1)
F   = np.maximum(F, 0.1)

RI = reflex_index(V, tau, C, F)

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
fig.suptitle('Figure 1 — Baseline RI Dynamics (T=150)', fontsize=13, fontweight='bold', color=LUX_BLUE)

ax = axes[0,0]
ax.plot(t, V, color=LUX_BLUE, lw=1.8, label='V(t) — volume')
ax.set_title('Information Volume V(t)'); ax.set_xlabel('t'); ax.legend()

ax = axes[0,1]
ax.plot(t, tau, color=LUX_GOLD, lw=1.8, label='τ(t) — toxicity')
ax.set_title('Toxicity τ(t)'); ax.set_xlabel('t'); ax.legend()

ax = axes[0,2]
ax.plot(t, C, color=LUX_GREEN, lw=1.8, label='C(t)')
ax.plot(t, F, color=LUX_GRAY,  lw=1.8, label='F(t)', linestyle='--')
ax.set_title('Absorption C(t) and Filtration F(t)'); ax.set_xlabel('t'); ax.legend()

ax = axes[1,0]
ax.plot(t, RI, color=LUX_RED, lw=2.2)
ax.axhline(RI.mean(), color=LUX_GOLD, lw=1.2, linestyle='--', label=f'Mean RI = {RI.mean():.3f}')
ax.fill_between(t, RI, alpha=0.12, color=LUX_RED)
ax.set_title('Reflexive Index RI(t)'); ax.set_xlabel('t'); ax.legend()

ax = axes[1,1]
ax.scatter(V*tau, C*F, c=t, cmap='viridis', s=18, alpha=0.7)
ax.set_xlabel('V·τ (numerator)'); ax.set_ylabel('C·F (denominator)')
ax.set_title('Phase portrait V·τ vs C·F')
sm = plt.cm.ScalarMappable(cmap='viridis', norm=plt.Normalize(0, T))
plt.colorbar(sm, ax=ax, label='t')

ax = axes[1,2]
ax.hist(RI, bins=25, color=LUX_BLUE, alpha=0.75, edgecolor='white')
ax.axvline(RI.mean(), color=LUX_GOLD, lw=1.5, linestyle='--', label=f'μ={RI.mean():.3f}')
ax.axvline(np.percentile(RI, 95), color=LUX_RED, lw=1.5, linestyle=':', label=f'P95={np.percentile(RI,95):.3f}')
ax.set_title('RI distribution'); ax.set_xlabel('RI'); ax.legend()

plt.tight_layout()
plt.savefig('figures/simulation_base.png', dpi=200, bbox_inches='tight')
plt.close()
print("✅ Figure 1 saved")

# ══════════════════════════════════════════════════════════════════════════════
# FIGURE 2 — Filtration scenarios
# ══════════════════════════════════════════════════════════════════════════════
fig, axes = plt.subplots(1, 3, figsize=(14, 5))
fig.suptitle('Figure 2 — Filtration F Scenarios', fontsize=13, fontweight='bold', color=LUX_BLUE)

scenarios = {
    'No filtration (F=0.1)':      lambda t: np.ones(len(t)) * 0.1,
    'Baseline F(t)':               lambda t: 1.5 + 0.15*t,
    'Reinforced F (+0.40/t)':      lambda t: 1.5 + 0.40*t,
}
colors = [LUX_RED, LUX_GRAY, LUX_GREEN]

V2   = 1.0 + 0.20*t
tau2 = 0.3 + 0.005*t
C2   = 2.0 + 0.10*t

for ax, (label, f_func), col in zip(axes, scenarios.items(), colors):
    F2 = np.maximum(f_func(t), 0.01)
    RI2 = reflex_index(V2, tau2, C2, F2)
    ax.plot(t, RI2, color=col, lw=2)
    ax.fill_between(t, RI2, alpha=0.12, color=col)
    ax.axhline(RI2.mean(), lw=1.2, linestyle='--', color=col, alpha=0.7,
               label=f'μ={RI2.mean():.3f}')
    ax.set_title(label, fontsize=10)
    ax.set_xlabel('t'); ax.set_ylabel('RI(t)')
    ax.legend(fontsize=9)

plt.tight_layout()
plt.savefig('figures/scenario_filtration.png', dpi=200, bbox_inches='tight')
plt.close()
print("✅ Figure 2 saved")

# ══════════════════════════════════════════════════════════════════════════════
# FIGURE 3 — Network topology and propagation
# ══════════════════════════════════════════════════════════════════════════════
np.random.seed(42)
n = 200

networks = {
    'Erdős–Rényi\n(random, p=0.03)':    nx.erdos_renyi_graph(n, 0.03, seed=42),
    'Barabási–Albert\n(scale-free, m=3)': nx.barabasi_albert_graph(n, 3, seed=42),
    'Watts–Strogatz\n(small-world)':      nx.watts_strogatz_graph(n, 6, 0.1, seed=42),
}

def simulate_propagation(G, T_prop=80, seed_frac=0.05, beta=0.15, gamma=0.05):
    np.random.seed(42)
    nodes = list(G.nodes())
    infected = set(np.random.choice(nodes, size=max(1, int(len(nodes)*seed_frac)), replace=False))
    recovered = set()
    history = [len(infected) / len(nodes)]
    for _ in range(T_prop):
        new_infected = set()
        new_recovered = set()
        for node in list(infected):
            if np.random.random() < gamma:
                new_recovered.add(node)
            else:
                for nb in G.neighbors(node):
                    if nb not in infected and nb not in recovered:
                        if np.random.random() < beta:
                            new_infected.add(nb)
        infected = (infected | new_infected) - new_recovered
        recovered |= new_recovered
        history.append(len(infected) / len(nodes))
    return history

fig, axes = plt.subplots(1, 3, figsize=(14, 5))
fig.suptitle('Figure 3 — Cognitive Contagion by Network Topology (n=200, β=0.15, γ=0.05)',
             fontsize=12, fontweight='bold', color=LUX_BLUE)

colors = [LUX_GRAY, LUX_RED, LUX_BLUE]
for ax, (label, G), col in zip(axes, networks.items(), colors):
    hist = simulate_propagation(G)
    degree_seq = sorted([d for _, d in G.degree()], reverse=True)
    ax.plot(hist, color=col, lw=2.2)
    ax.fill_between(range(len(hist)), hist, alpha=0.12, color=col)
    ax.set_title(f'{label}\nMax degree={max(degree_seq)}, Peak={max(hist):.2f}',
                 fontsize=9)
    ax.set_xlabel('Time steps'); ax.set_ylabel('Fraction infected')
    ax.set_ylim(0, 1)

plt.tight_layout()
plt.savefig('figures/propagation_reseau.png', dpi=200, bbox_inches='tight')
plt.close()
print("✅ Figure 3 saved")

# ══════════════════════════════════════════════════════════════════════════════
# FIGURE 4 — Doctrinal coherence D oscillation
# ══════════════════════════════════════════════════════════════════════════════
T_doc = 200
t_doc = np.arange(T_doc)

threat = 0.3 + 0.3*np.sin(2*np.pi*t_doc/50) + 0.1*np.random.randn(T_doc)
threat = np.clip(threat, 0, 1)

D = np.zeros(T_doc)
D[0] = 0.2
alpha_raise = 0.08
alpha_lower = 0.05
threshold   = 0.5

for i in range(1, T_doc):
    if threat[i] > threshold:
        D[i] = min(1.0, D[i-1] + alpha_raise * (threat[i] - threshold))
    else:
        D[i] = max(0.05, D[i-1] - alpha_lower * (threshold - threat[i]))

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 7), sharex=True)
fig.suptitle('Figure 4 — Doctrinal Coherence D Oscillation\n(Auftragstaktik regime D→1 / Iga regime D→0)',
             fontsize=12, fontweight='bold', color=LUX_BLUE)

ax1.plot(t_doc, threat, color=LUX_RED, lw=1.8, label='Threat level')
ax1.axhline(threshold, color=LUX_GOLD, lw=1.2, linestyle='--', label=f'Threshold = {threshold}')
ax1.fill_between(t_doc, threat, threshold, where=threat>threshold,
                 alpha=0.2, color=LUX_RED, label='Auftragstaktik zone')
ax1.fill_between(t_doc, threat, threshold, where=threat<=threshold,
                 alpha=0.15, color=LUX_BLUE, label='Iga zone')
ax1.set_ylabel('Threat level'); ax1.legend(fontsize=9)

ax2.plot(t_doc, D, color=LUX_BLUE, lw=2.2, label='D(t)')
ax2.axhline(0.5, color=LUX_GRAY, lw=0.8, linestyle=':')
ax2.fill_between(t_doc, D, 0.5, where=D>0.5, alpha=0.15, color=LUX_BLUE, label='Auftragstaktik')
ax2.fill_between(t_doc, D, 0.5, where=D<=0.5, alpha=0.12, color=LUX_GREEN, label='Iga')
ax2.set_xlabel('t'); ax2.set_ylabel('D(t)'); ax2.set_ylim(0, 1)
ax2.legend(fontsize=9)

plt.tight_layout()
plt.savefig('figures/doctrinal_oscillation.png', dpi=200, bbox_inches='tight')
plt.close()
print("✅ Figure 4 saved")

# ══════════════════════════════════════════════════════════════════════════════
# FIGURE 5 — Sensitivity analysis C×F
# ══════════════════════════════════════════════════════════════════════════════
C_vals = np.linspace(0.5, 5.0, 40)
F_vals = np.linspace(0.5, 5.0, 40)
CC, FF = np.meshgrid(C_vals, F_vals)

V_fixed   = 3.0
tau_fixed = 0.5
RI_grid = reflex_index(V_fixed, tau_fixed, CC, FF)

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle(f'Figure 5 — Sensitivity Analysis: RI as a function of C and F\n(V={V_fixed}, τ={tau_fixed})',
             fontsize=12, fontweight='bold', color=LUX_BLUE)

im = axes[0].contourf(C_vals, F_vals, RI_grid, levels=25, cmap='RdYlGn_r')
axes[0].contour(C_vals, F_vals, RI_grid, levels=[0.3, 0.6, 1.0], colors='black',
                linewidths=0.8, linestyles='--')
plt.colorbar(im, ax=axes[0], label='RI')
axes[0].set_xlabel('C — cognitive absorption'); axes[0].set_ylabel('F — algorithmic filtration')
axes[0].set_title('RI heatmap\n(green = low risk, red = high risk)')

CF_product = (CC * FF).flatten()
RI_flat = RI_grid.flatten()
sort_idx = np.argsort(CF_product)
axes[1].scatter(CF_product[sort_idx], RI_flat[sort_idx], c=RI_flat[sort_idx],
                cmap='RdYlGn_r', s=12, alpha=0.6)
axes[1].set_xlabel('C × F (combined defence capacity)')
axes[1].set_ylabel('RI')
axes[1].set_title('RI vs C×F product\n(RI stable only when C×F is high)')
axes[1].axvline(4.0, color=LUX_GOLD, lw=1.2, linestyle='--', label='C×F = 4 threshold')
axes[1].legend()

plt.tight_layout()
plt.savefig('figures/sensibilite_CF.png', dpi=200, bbox_inches='tight')
plt.close()
print("✅ Figure 5 saved")

# ══════════════════════════════════════════════════════════════════════════════
# FIGURE 6 — Coupled feedback loop (v3.0 correction)
# τ(t+1) = τ(t) + α·RI(t)   [RI amplifies toxicity]
# C(t+1) = C(t) - β·RI(t)   [RI erodes cognitive resilience]
# ══════════════════════════════════════════════════════════════════════════════
def simulate_coupled(T_c=150, alpha=0.015, beta=0.008,
                     V0=1.0, tau0=0.3, C0=3.0, F0=1.5,
                     V_rate=0.12, F_rate=0.10):
    """Coupled dynamic system with RI feedback into tau and C."""
    np.random.seed(42)
    V_c   = np.zeros(T_c)
    tau_c = np.zeros(T_c)
    C_c   = np.zeros(T_c)
    F_c   = np.zeros(T_c)
    RI_c  = np.zeros(T_c)

    V_c[0]   = V0
    tau_c[0] = tau0
    C_c[0]   = C0
    F_c[0]   = F0
    RI_c[0]  = reflex_index(V_c[0], tau_c[0], C_c[0], F_c[0])

    noise = np.random.randn(T_c) * 0.01

    for i in range(1, T_c):
        V_c[i]   = V_c[i-1]   + V_rate  + noise[i]
        F_c[i]   = F_c[i-1]   + F_rate  + noise[i]*0.5
        # Feedback: RI degrades C and amplifies tau
        tau_c[i] = np.clip(tau_c[i-1] + alpha * RI_c[i-1] + noise[i]*0.5, 0.01, 2.0)
        C_c[i]   = max(0.05, C_c[i-1] - beta  * RI_c[i-1] + noise[i]*0.3)
        RI_c[i]  = reflex_index(V_c[i], tau_c[i], C_c[i], F_c[i])

    return V_c, tau_c, C_c, F_c, RI_c

# Simulate: with and without feedback
V_nf = 1.0 + 0.12*np.arange(T) + np.random.RandomState(42).randn(T)*0.01
tau_nf = np.ones(T) * 0.3
C_nf = 3.0 + 0.10*np.arange(T)
F_nf = 1.5 + 0.10*np.arange(T)
RI_nf = reflex_index(V_nf, tau_nf, C_nf, F_nf)

V_fb, tau_fb, C_fb, F_fb, RI_fb = simulate_coupled()

fig, axes = plt.subplots(2, 3, figsize=(15, 9))
fig.suptitle('Figure 6 — Coupled Feedback Dynamics vs Open-Loop\n'
             'τ(t+1) = τ(t) + α·RI(t),   C(t+1) = C(t) − β·RI(t)',
             fontsize=12, fontweight='bold', color=LUX_BLUE)
t_c = np.arange(T)

# RI comparison
ax = axes[0,0]
ax.plot(t_c, RI_nf, color=LUX_GRAY, lw=1.8, linestyle='--', label='Open-loop (exogenous)')
ax.plot(t_c, RI_fb, color=LUX_RED,  lw=2.2, label='Coupled feedback')
ax.fill_between(t_c, RI_fb, RI_nf, alpha=0.15, color=LUX_RED)
ax.set_title('RI(t): open-loop vs feedback'); ax.set_xlabel('t'); ax.legend()

# Toxicity amplification
ax = axes[0,1]
ax.plot(t_c, tau_nf, color=LUX_GRAY, lw=1.8, linestyle='--', label='τ open-loop (const)')
ax.plot(t_c, tau_fb, color=LUX_GOLD, lw=2.2, label='τ coupled (feedback)')
ax.set_title('Toxicity τ(t) — RI amplification'); ax.set_xlabel('t'); ax.legend()

# Cognitive resilience erosion
ax = axes[0,2]
ax.plot(t_c, C_nf, color=LUX_GRAY,  lw=1.8, linestyle='--', label='C open-loop')
ax.plot(t_c, C_fb, color=LUX_GREEN, lw=2.2, label='C coupled (eroded by RI)')
ax.set_title('Cognitive resilience C(t) — RI erosion'); ax.set_xlabel('t'); ax.legend()

# Phase portrait tau vs C
ax = axes[1,0]
sc = ax.scatter(tau_fb, C_fb, c=t_c, cmap='plasma', s=18, alpha=0.8)
ax.set_xlabel('τ(t)'); ax.set_ylabel('C(t)')
ax.set_title('Phase portrait τ vs C\n(coupled system)')
plt.colorbar(sc, ax=ax, label='t')

# RI trajectory with bifurcation marker
ax = axes[1,1]
ax.plot(t_c, RI_fb, color=LUX_RED, lw=2)
# Find approximate bifurcation (RI starts accelerating)
dRI = np.gradient(RI_fb)
bifurc = np.where(dRI > dRI.mean() + dRI.std())[0]
if len(bifurc):
    ax.axvline(bifurc[0], color=LUX_GOLD, lw=1.5, linestyle='--',
               label=f'Acceleration onset t≈{bifurc[0]}')
ax.set_title('RI(t) — coupled, with acceleration marker')
ax.set_xlabel('t'); ax.legend()

# Sensitivity: alpha and beta effect on final RI
alphas = np.linspace(0.001, 0.04, 15)
betas  = np.linspace(0.001, 0.02, 15)
AA, BB = np.meshgrid(alphas, betas)
final_RI = np.zeros_like(AA)
for i in range(AA.shape[0]):
    for j in range(AA.shape[1]):
        _, _, _, _, ri = simulate_coupled(alpha=AA[i,j], beta=BB[i,j])
        final_RI[i,j] = ri[-1]

ax = axes[1,2]
im = ax.contourf(alphas, betas, final_RI, levels=20, cmap='RdYlGn_r')
plt.colorbar(im, ax=ax, label='Final RI')
ax.set_xlabel('α (toxicity amplification)'); ax.set_ylabel('β (resilience erosion)')
ax.set_title('Final RI sensitivity\nto feedback coefficients α, β')

plt.tight_layout()
plt.savefig('figures/feedback_coupled.png', dpi=200, bbox_inches='tight')
plt.close()
print("✅ Figure 6 saved")

# ══════════════════════════════════════════════════════════════════════════════
# CSV export
# ══════════════════════════════════════════════════════════════════════════════
df = pd.DataFrame({'t': t, 'V': V, 'tau': tau, 'C': C, 'F': F, 'RI': RI})
df.to_csv('simulation_data.csv', index=False)

print("\n📊 Descriptive statistics — RI(t):")
print(df['RI'].describe().round(4))
print("\n✅ All figures and CSV exported successfully.")