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/hcl_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)

    # Step 1: 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()

    # Step 2: Group by Temperature_C and HCl_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', 'HCl_M'])[recovery_cols].mean().reset_index()

    # Step 3: Compute Selectivity_Index
    grouped['Selectivity_Index'] = (grouped['Nd_Recovery_pct'] + grouped['Dy_Recovery_pct'] + grouped['Pr_Recovery_pct']) / (3.0 * grouped['Fe_Recovery_pct'])

    # Step 4: Compute Weighted_Econ_Recovery
    # Weights: Nd: price=120 $/kg, wt%=24.5; Dy: price=350, wt%=1.2; Pr: price=95, wt%=0.85
    nd_weight = 120 * 24.5
    dy_weight = 350 * 1.2
    pr_weight = 95 * 0.85
    total_weight = nd_weight + dy_weight + pr_weight
    
    grouped['Weighted_Econ_Recovery'] = (
        grouped['Nd_Recovery_pct'] * nd_weight +
        grouped['Dy_Recovery_pct'] * dy_weight +
        grouped['Pr_Recovery_pct'] * pr_weight
    ) / total_weight

    # Step 5: Rename HCl_M to Acid_Concentration_M
    grouped = grouped.rename(columns={'HCl_M': 'Acid_Concentration_M'})

    # Step 6: Add Acid_Type column
    grouped['Acid_Type'] = 'HCl'

    # 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
    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(2, 2, figsize=(14, 10))
    fig.suptitle('HCl Leaching Screen: Recovery & Selectivity vs Temperature and Concentration', fontsize=13)

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

    # Plot 1: Nd Recovery
    ax = axes[0, 0]
    for i, conc in enumerate(concentrations):
        subset = result_df[result_df['Acid_Concentration_M'] == conc]
        ax.plot(subset['Temperature_C'], subset['Nd_Recovery_pct'], 'o-', color=colors[i], label=f'{conc} M')
    ax.set_xlabel('Temperature (°C)')
    ax.set_ylabel('Nd Recovery (%)')
    ax.set_title('Nd Recovery')
    ax.legend(title='[HCl]', fontsize=8)
    ax.grid(True, alpha=0.3)

    # Plot 2: Fe Recovery
    ax = axes[0, 1]
    for i, conc in enumerate(concentrations):
        subset = result_df[result_df['Acid_Concentration_M'] == conc]
        ax.plot(subset['Temperature_C'], subset['Fe_Recovery_pct'], 's-', color=colors[i], label=f'{conc} M')
    ax.set_xlabel('Temperature (°C)')
    ax.set_ylabel('Fe Recovery (%)')
    ax.set_title('Fe Recovery (co-dissolution)')
    ax.legend(title='[HCl]', fontsize=8)
    ax.grid(True, alpha=0.3)

    # Plot 3: Selectivity Index
    ax = axes[1, 0]
    for i, conc in enumerate(concentrations):
        subset = result_df[result_df['Acid_Concentration_M'] == conc]
        ax.plot(subset['Temperature_C'], subset['Selectivity_Index'], 'D-', color=colors[i], label=f'{conc} M')
    ax.set_xlabel('Temperature (°C)')
    ax.set_ylabel('Selectivity Index (REE/Fe)')
    ax.set_title('Selectivity Index')
    ax.legend(title='[HCl]', fontsize=8)
    ax.axhline(y=1.0, color='red', linestyle='--', alpha=0.5, label='S=1')
    ax.grid(True, alpha=0.3)

    # Plot 4: Weighted Economic Recovery
    ax = axes[1, 1]
    for i, conc in enumerate(concentrations):
        subset = result_df[result_df['Acid_Concentration_M'] == conc]
        ax.plot(subset['Temperature_C'], subset['Weighted_Econ_Recovery'], '^-', color=colors[i], label=f'{conc} M')
    ax.set_xlabel('Temperature (°C)')
    ax.set_ylabel('Weighted Econ Recovery (%)')
    ax.set_title('Weighted Economic Recovery')
    ax.legend(title='[HCl]', fontsize=8)
    ax.grid(True, alpha=0.3)

    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))
