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_160016/data/hcl_plate1_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)

    # Calculate REE_Selectivity_Index
    # SI = mean(Nd_Recovery_pct, Dy_Recovery_pct, Pr_Recovery_pct, Co_Recovery_pct) / (Fe_Recovery_pct + 1)
    df['REE_Selectivity_Index'] = (
        (df['Nd_Recovery_pct'] + df['Dy_Recovery_pct'] + df['Pr_Recovery_pct'] + df['Co_Recovery_pct']) / 4.0
    ) / (df['Fe_Recovery_pct'] + 1.0)

    # Calculate Economic_Recovery_Value
    # Value = (Nd_Recovery_pct/100 * 245 * 120) + (Dy_Recovery_pct/100 * 12 * 350) + 
    #         (Pr_Recovery_pct/100 * 8.5 * 95) + (Co_Recovery_pct/100 * 3.2 * 35)
    df['Economic_Recovery_Value'] = (
        (df['Nd_Recovery_pct'] / 100.0 * 245 * 120) +
        (df['Dy_Recovery_pct'] / 100.0 * 12 * 350) +
        (df['Pr_Recovery_pct'] / 100.0 * 8.5 * 95) +
        (df['Co_Recovery_pct'] / 100.0 * 3.2 * 35)
    )

    # Build metrics list with exact column names
    metrics_list = []
    for _, row in df.iterrows():
        metrics_list.append({
            'Temperature_C': float(row['Temperature_C']),
            'HCl_M': float(row['HCl_M']),
            'REE_Selectivity_Index': round(float(row['REE_Selectivity_Index']), 4),
            'Economic_Recovery_Value': round(float(row['Economic_Recovery_Value']), 2)
        })

    # Create plot
    output_dir = Path("/Users/maxim.ziatdinov/Code/SciLink/planning_session_20260319_160016/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, 5))

    # Get unique temperatures and concentrations for heatmap
    temps = sorted(df['Temperature_C'].unique())
    concs = sorted(df['HCl_M'].unique())

    # Create pivot tables
    si_pivot = df.pivot_table(index='Temperature_C', columns='HCl_M', values='REE_Selectivity_Index')
    ev_pivot = df.pivot_table(index='Temperature_C', columns='HCl_M', values='Economic_Recovery_Value')

    # Plot REE_Selectivity_Index heatmap
    im1 = axes[0].imshow(si_pivot.values, aspect='auto', origin='lower', cmap='viridis')
    axes[0].set_xticks(range(len(concs)))
    axes[0].set_xticklabels([f'{c}' for c in concs], rotation=45, fontsize=7)
    axes[0].set_yticks(range(len(temps)))
    axes[0].set_yticklabels([f'{t}' for t in temps])
    axes[0].set_xlabel('HCl (M)')
    axes[0].set_ylabel('Temperature (°C)')
    axes[0].set_title('REE Selectivity Index\n(higher = better Fe rejection)')
    plt.colorbar(im1, ax=axes[0], shrink=0.8)

    # Annotate with values
    for i in range(len(temps)):
        for j in range(len(concs)):
            if i < si_pivot.values.shape[0] and j < si_pivot.values.shape[1]:
                val = si_pivot.values[i, j]
                if not np.isnan(val):
                    axes[0].text(j, i, f'{val:.2f}', ha='center', va='center', fontsize=5, color='white' if val < si_pivot.values[~np.isnan(si_pivot.values)].mean() else 'black')

    # Plot Economic_Recovery_Value heatmap
    im2 = axes[1].imshow(ev_pivot.values, aspect='auto', origin='lower', cmap='YlOrRd')
    axes[1].set_xticks(range(len(concs)))
    axes[1].set_xticklabels([f'{c}' for c in concs], rotation=45, fontsize=7)
    axes[1].set_yticks(range(len(temps)))
    axes[1].set_yticklabels([f'{t}' for t in temps])
    axes[1].set_xlabel('HCl (M)')
    axes[1].set_ylabel('Temperature (°C)')
    axes[1].set_title('Economic Recovery Value ($/tonne)\n(higher = more value recovered)')
    plt.colorbar(im2, ax=axes[1], shrink=0.8)

    # Annotate with values
    for i in range(len(temps)):
        for j in range(len(concs)):
            if i < ev_pivot.values.shape[0] and j < ev_pivot.values.shape[1]:
                val = ev_pivot.values[i, j]
                if not np.isnan(val):
                    axes[1].text(j, i, f'{val:.0f}', ha='center', va='center', fontsize=5, color='white' if val > ev_pivot.values[~np.isnan(ev_pivot.values)].mean() else 'black')

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

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