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

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

    # Define grouping columns (inputs)
    group_cols = ['Pretreatment', 'Roast_Temp_C', 'Lixiviant', 'Acid_Concentration_M']

    # Recovery columns
    recovery_cols = ['Nd_Recovery_pct', 'Dy_Recovery_pct', 'Pr_Recovery_pct', 'Fe_Recovery_pct',
                     'Co_Recovery_pct', 'B_Recovery_pct', 'Ni_Recovery_pct']

    # Group by input parameters and compute means across duplicates
    grouped = df.groupby(group_cols, as_index=False)[recovery_cols].mean()

    # Compute Selectivity_Index = (mean_Nd + mean_Dy + mean_Pr) / (3 * mean_Fe)
    # Handle division by zero
    grouped['Selectivity_Index'] = np.where(
        grouped['Fe_Recovery_pct'] > 0,
        (grouped['Nd_Recovery_pct'] + grouped['Dy_Recovery_pct'] + grouped['Pr_Recovery_pct']) / (3.0 * grouped['Fe_Recovery_pct']),
        np.nan
    )

    # Compute Weighted_Econ_Recovery
    # Weights: Nd: price=120, wt%=24.5; Dy: price=350, wt%=1.2; Pr: price=95, wt%=0.85
    w_Nd = 120 * 24.5  # 2940
    w_Dy = 350 * 1.2   # 420
    w_Pr = 95 * 0.85   # 80.75
    w_total = w_Nd + w_Dy + w_Pr  # 3440.75

    grouped['Weighted_Econ_Recovery'] = (
        grouped['Nd_Recovery_pct'] * w_Nd +
        grouped['Dy_Recovery_pct'] * w_Dy +
        grouped['Pr_Recovery_pct'] * w_Pr
    ) / w_total

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

    # Build metrics list - one row per group
    output_cols = group_cols + ['Nd_Recovery_pct', 'Dy_Recovery_pct', 'Pr_Recovery_pct',
                                'Fe_Recovery_pct', 'Selectivity_Index', 'Weighted_Econ_Recovery']
    metrics_list = grouped[output_cols].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=(16, 8))

    # Subplot 1: Selectivity Index by condition
    # Create a label for each group
    grouped['label'] = (
        grouped['Pretreatment'].astype(str) + '\n' +
        grouped['Roast_Temp_C'].astype(str) + '°C\n' +
        grouped['Lixiviant'].astype(str) + '\n' +
        grouped['Acid_Concentration_M'].astype(str) + 'M'
    )

    # Filter H2SO4 only for cleaner plot, sort by SI
    h2so4_data = grouped[grouped['Lixiviant'] == 'H2SO4'].copy()
    acetic_data = grouped[grouped['Lixiviant'] == 'AceticAcid'].copy()

    # Plot SI vs Nd_Recovery colored by roast temp for H2SO4
    ax1 = axes[0]
    temps = h2so4_data['Roast_Temp_C'].unique()
    colors_temp = {400: '#2196F3', 500: '#4CAF50', 600: '#FF9800', 800: '#F44336'}
    markers_pre = {'NaOH': 'o', 'Direct': 's'}

    for pre in ['NaOH', 'Direct']:
        for temp in sorted(temps):
            mask = (h2so4_data['Pretreatment'] == pre) & (h2so4_data['Roast_Temp_C'] == temp)
            subset = h2so4_data[mask]
            if len(subset) > 0:
                ax1.scatter(
                    subset['Nd_Recovery_pct'], subset['Selectivity_Index'],
                    c=colors_temp.get(temp, 'gray'),
                    marker=markers_pre.get(pre, 'o'),
                    s=100, edgecolors='black', linewidth=0.5,
                    label=f'{pre} {temp}°C',
                    alpha=0.85
                )

    # Add acetic acid points
    for pre in ['NaOH', 'Direct']:
        for temp in sorted(acetic_data['Roast_Temp_C'].unique()):
            mask = (acetic_data['Pretreatment'] == pre) & (acetic_data['Roast_Temp_C'] == temp)
            subset = acetic_data[mask]
            if len(subset) > 0:
                ax1.scatter(
                    subset['Nd_Recovery_pct'], subset['Selectivity_Index'],
                    c=colors_temp.get(temp, 'gray'),
                    marker=markers_pre.get(pre, 'o'),
                    s=100, edgecolors='red', linewidth=2.0,
                    alpha=0.85
                )

    ax1.axhline(y=10, color='green', linestyle='--', alpha=0.7, label='SI=10 target')
    ax1.axvline(x=80, color='blue', linestyle='--', alpha=0.5, label='80% Nd target')
    ax1.set_xlabel('Nd Recovery (%)', fontsize=12)
    ax1.set_ylabel('Selectivity Index (REE/Fe)', fontsize=12)
    ax1.set_title('Selectivity vs Nd Recovery\n(circles=NaOH, squares=Direct; red edge=AceticAcid)', fontsize=11)
    ax1.legend(fontsize=7, loc='upper left', ncol=2)
    ax1.set_xlim(0, 105)
    ax1.grid(True, alpha=0.3)

    # Subplot 2: Bar chart of Weighted_Econ_Recovery for top conditions
    ax2 = axes[1]
    top_conditions = grouped.nlargest(15, 'Weighted_Econ_Recovery').copy()
    top_conditions['short_label'] = (
        top_conditions['Pretreatment'].str[:4] + ' ' +
        top_conditions['Roast_Temp_C'].astype(str) + '°C ' +
        top_conditions['Lixiviant'].str[:5] + ' ' +
        top_conditions['Acid_Concentration_M'].astype(str) + 'M'
    )
    bar_colors = ['#4CAF50' if si > 10 else '#FF9800' if si > 5 else '#F44336'
                  for si in top_conditions['Selectivity_Index']]
    bars = ax2.barh(range(len(top_conditions)), top_conditions['Weighted_Econ_Recovery'],
                    color=bar_colors, edgecolor='black', linewidth=0.5)
    ax2.set_yticks(range(len(top_conditions)))
    ax2.set_yticklabels(top_conditions['short_label'], fontsize=8)
    ax2.set_xlabel('Weighted Econ Recovery (%)', fontsize=12)
    ax2.set_title('Top 15 Conditions by Weighted Econ Recovery\n(Green=SI>10, Orange=SI>5, Red=SI<5)', fontsize=11)
    ax2.invert_yaxis()
    ax2.grid(True, alpha=0.3, axis='x')

    # Add SI values as text
    for i, (_, row) in enumerate(top_conditions.iterrows()):
        ax2.text(row['Weighted_Econ_Recovery'] + 0.5, i,
                 f'SI={row["Selectivity_Index"]:.1f}', va='center', fontsize=7)

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

    # Output results
    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))
