import xarray as xr
import matplotlib.pyplot as plt
import numpy as np

# 1. Caricamento dataset
ds03 = xr.open_dataset('tau_03/04008.atmos_daily.nc', decode_times=False)
ds00 = xr.open_dataset('../00/04008.atmos_daily.nc', decode_times=False)

def get_data(ds):
    # Selezione Meridiani Planum (Lat 0, Lon 358)
    p = ds.sel(lat=0.0, lon=358.0, method='nearest')
    ls = p.areo.values[:, 0] % 360
    
    # Ordinamento per Ls
    idx = np.argsort(ls)
    ls, u, v, tau = ls[idx], p.ukd.values[idx], p.vkd.values[idx], p.stress.values[idx]
    
    # Calcolo componenti stress
    ws = np.sqrt(u**2 + v**2)
    ws_safe = np.where(ws > 0, ws, 1.0)
    tx = np.where(ws > 0, tau * (u / ws_safe), 0) # Est-Ovest
    ty = np.where(ws > 0, tau * (v / ws_safe), 0) # Nord-Sud
    
    return ls, tx, ty, tau, u, v

ls03, tx03, ty03, ttot03, u03, v03 = get_data(ds03)
ls00, tx00, ty00, ttot00, u00, v00 = get_data(ds00)

# 2. CALCOLO RDD E DEBUG FISICO
soglia = 0.0225
mask03 = ttot03 > soglia

print("--- ANALISI FISICA CASE 03 (400ka) ---")
if np.any(mask03):
    m_tx, m_ty = tx03[mask03].mean(), ty03[mask03].mean()
    rdd = (np.degrees(np.arctan2(m_tx, m_ty)) + 360) % 360
    print(f"RDD Calcolato: {rdd:.2f}° (Fenton: 317°)")
    print(f"Stress X (Ovest): {m_tx:.5f} | Stress Y (Nord): {m_ty:.5f}")
    print(f"Rapporto Y/X: {abs(m_ty/m_tx):.3f} (Deve essere ~0.9 per i 317°)")
else:
    print("Nessun valore sopra soglia!")

# 3. VERIFICA GRADIENTE TERMICO (Perché il vento non va a Nord?)
# Prendiamo la temperatura superficiale (ts) media al Polo Nord e all'Equatore
ts_equatore = ds03.sel(lat=0, method='nearest').ts.mean().values
ts_polonord = ds03.sel(lat=80, method='nearest').ts.mean().values
print(f"\nTemp. Media Equatore: {ts_equatore:.1f} K")
print(f"Temp. Media Polo Nord: {ts_polonord:.1f} K")
print(f"Differenza T (Nord - Eq): {ts_polonord - ts_equatore:.1f} K")

# 4. GRAFICO FINALE
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10), sharex=True)

# Stress Meridionale (Nord-Sud)
ax1.plot(ls00, ty00, label='Oggi (00)', color='blue', alpha=0.4)
ax1.plot(ls03, ty03, label='400ka (03)', color='red', lw=1.5)
ax1.axhline(0, color='black', lw=0.8)
ax1.axhline(soglia, color='green', ls='--', alpha=0.5, label='Soglia Fenton')
ax1.set_ylabel("Tau_y [Pa] (Positivo = NORD)")
ax1.set_title(f"Componente Meridionale - RDD: {rdd:.1f}°")
ax1.legend()

# Stress Zonale (Est-Ovest)
ax2.plot(ls00, tx00, label='Oggi (00)', color='blue', alpha=0.4)
ax2.plot(ls03, tx03, label='400ka (03)', color='darkred', lw=1.5)
ax2.axhline(0, color='black', lw=0.8)
ax2.set_ylabel("Tau_x [Pa] (Negativo = OVEST)")
ax2.set_xlabel("Ls [Gradi]")

plt.tight_layout()
plt.savefig("2debug_rdd_fenton.png", dpi=300)
plt.show()