import sys
import json
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from pathlib import Path

try:
    # Accept file path as command-line argument
    if len(sys.argv) > 1:
        data_path = sys.argv[1]
    else:
        data_path = "/Users/maxim.ziatdinov/Code/SciLink/planning_session_20260319_194604/data/h2so4_plate_icpms_results.csv"

    # Read data
    df = pd.read_csv(data_path)

    # Load sidecar JSON if exists
    sidecar_path = Path(data_path).with_suffix('.json')
    sidecar_conditions = {}
    if sidecar_path.exists():
        with open(sidecar_path) as f:
            sidecar_conditions = json.load(f)

    # Filter to experimental rows only (A-F, exclude G and H blanks)
    experimental_rows = ['A', 'B', 'C', 'D', 'E', 'F']
    df_exp = df[df['Row'].isin(experimental_rows)].copy()

    # Group by Temperature_C and H2SO4_M, compute mean across triplicates
    recovery_cols = ['Nd_Recovery_pct', 'Dy_Recovery_pct', 'Pr_Recovery_pct',
                     'Fe_Recovery_pct', 'Co_Recovery_pct', 'B_Recovery_pct',
                     'Cu_Recovery_pct', 'Al_Recovery_pct', 'Ni_Recovery_pct']

    grouped = df_exp.groupby(['Temperature_C', 'H2SO4_M'])[recovery_cols].mean().reset_index()

    # Compute Selectivity_Index = (mean_Nd + mean_Dy + mean_Pr) / (3 * mean_Fe)
    grouped['Selectivity_Index'] = (
        (grouped['Nd_Recovery_pct'] + grouped['Dy_Recovery_pct'] + grouped['Pr_Recovery_pct']) /
        (3.0 * grouped['Fe_Recovery_pct'])
    )

    # Compute Weighted_Econ_Recovery
    # Weights: Nd: 120 $/kg * 24.5 wt%, Dy: 350 $/kg * 1.2 wt%, Pr: 95 $/kg * 0.85 wt%
    w_nd = 120.0 * 24.5
    w_dy = 350.0 * 1.2
    w_pr = 95.0 * 0.85
    w_total = w_nd + w_dy + w_pr

    grouped['Weighted_Econ_Recovery'] = (
        grouped['Nd_Recovery_pct'] * w_nd +
        grouped['Dy_Recovery_pct'] * w_dy +
        grouped['Pr_Recovery_pct'] * w_pr
    ) / w_total

    # Rename H2SO4_M to Acid_Concentration_M
    grouped = grouped.rename(columns={'H2SO4_M': 'Acid_Concentration_M'})

    # Add Acid_Type column
    grouped['Acid_Type'] = 'H2SO4'

    # Select output columns
    output_cols = ['Temperature_C', 'Acid_Concentration_M', 'Acid_Type',
                   'Nd_Recovery_pct', 'Dy_Recovery_pct', 'Pr_Recovery_pct',
                   'Fe_Recovery_pct', 'Selectivity_Index', 'Weighted_Econ_Recovery']
    result_df = grouped[output_cols].copy()

    # Round for cleanliness
    for col in ['Nd_Recovery_pct', 'Dy_Recovery_pct', 'Pr_Recovery_pct',
                'Fe_Recovery_pct', 'Selectivity_Index', 'Weighted_Econ_Recovery']:
        result_df[col] = result_df[col].round(4)

    # Convert to list of dicts for multi-row output
    metrics_list = result_df.to_dict(orient='records')

    # ---- PLOTTING ----
    output_dir = Path("/Users/maxim.ziatdinov/Code/SciLink/planning_session_20260319_194604/scalarizer_outputs")
    output_dir.mkdir(parents=True, exist_ok=True)
    plot_path = output_dir / f"debug_{Path(data_path).stem}.png"

    fig, axes = plt.subplots(1, 2, figsize=(14, 6))

    # Get unique concentrations and temperatures
    concentrations = sorted(result_df['Acid_Concentration_M'].unique())
    temperatures = sorted(result_df['Temperature_C'].unique())
    colors = plt.cm.viridis(np.linspace(0, 1, len(concentrations)))

    # Subplot 1: Nd Recovery and Fe Recovery vs Temperature for each concentration
    ax1 = axes[0]
    for i, conc in enumerate(concentrations):
        subset = result_df[result_df['Acid_Concentration_M'] == conc]
        ax1.plot(subset['Temperature_C'], subset['Nd_Recovery_pct'], 'o-',
                 color=colors[i], label=f'{conc} M - Nd', linewidth=2, markersize=6)
        ax1.plot(subset['Temperature_C'], subset['Fe_Recovery_pct'], 's--',
                 color=colors[i], alpha=0.5, linewidth=1.5, markersize=5)
    ax1.set_xlabel('Temperature (°C)', fontsize=12)
    ax1.set_ylabel('Recovery (%)', fontsize=12)
    ax1.set_title('H₂SO₄ Leaching: Nd (solid) & Fe (dashed) Recovery', fontsize=11)
    ax1.legend(fontsize=8, title='[H₂SO₄]', ncol=2)
    ax1.grid(True, alpha=0.3)

    # Subplot 2: Selectivity Index vs Temperature for each concentration
    ax2 = axes[1]
    for i, conc in enumerate(concentrations):
        subset = result_df[result_df['Acid_Concentration_M'] == conc]
        ax2.plot(subset['Temperature_C'], subset['Selectivity_Index'], 'o-',
                 color=colors[i], label=f'{conc} M', linewidth=2, markersize=6)
    ax2.axhline(y=3.0, color='red', linestyle=':', alpha=0.7, label='Target S=3')
    ax2.set_xlabel('Temperature (°C)', fontsize=12)
    ax2.set_ylabel('Selectivity Index (REE/Fe)', fontsize=12)
    ax2.set_title('H₂SO₄ Leaching: REE/Fe Selectivity Index', fontsize=11)
    ax2.legend(fontsize=9, title='[H₂SO₄]')
    ax2.grid(True, alpha=0.3)

    # Overall title with key stats
    best_idx = result_df['Selectivity_Index'].idxmax()
    best_row = result_df.loc[best_idx]
    fig.suptitle(
        f"H₂SO₄ Plate: Best Selectivity = {best_row['Selectivity_Index']:.2f} "
        f"at {best_row['Temperature_C']}°C, {best_row['Acid_Concentration_M']} M\n"
        f"(Nd={best_row['Nd_Recovery_pct']:.1f}%, Fe={best_row['Fe_Recovery_pct']:.1f}%)",
        fontsize=12, fontweight='bold', y=1.02
    )

    plt.tight_layout()
    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
    plt.close()

    # Output JSON
    result = {
        "metrics": metrics_list,
        "plot_path": str(plot_path)
    }
    print(json.dumps(result))

except Exception as e:
    import traceback
    error_result = {
        "metrics": None,
        "plot_path": None,
        "error": str(e),
        "traceback": traceback.format_exc()
    }
    print(json.dumps(error_result))
