"""
Visualize AFRES v2 results.
"""
import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

def load(path="/app/afres_v2_results.json"):
    with open(path) as f:
        return json.load(f)

def plot_search_history(results, save="/app/v2_search_history.png"):
    hist = results.get("search_history", [])
    if not hist:
        return
    iters = [h["iteration"] for h in hist]
    fitnesses = [h.get("fitness", 0) for h in hist]
    maes = [h.get("mean_mae", 0) for h in hist]
    ics = [h.get("mean_ic", 0) for h in hist]

    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 5))

    ax1.scatter(iters, fitnesses, c='steelblue', alpha=0.6, s=30)
    ax1.plot(sorted(set(iters)), [max(f for i, f in zip(iters, fitnesses) if i == ii)
                                    for ii in sorted(set(iters))], 'r-', linewidth=2)
    ax1.set_xlabel("Iteration")
    ax1.set_ylabel("Fitness (Mean IC)")
    ax1.set_title("Rubric Search: Fitness Landscape")
    ax1.grid(True, alpha=0.3)

    ax2.scatter(iters, ics, c='green', alpha=0.6, s=30)
    ax2.set_xlabel("Iteration")
    ax2.set_ylabel("Mean IC")
    ax2.set_title("Mean IC per Evaluated Rubric")
    ax2.grid(True, alpha=0.3)

    ax3.scatter(iters, maes, c='orange', alpha=0.6, s=30)
    ax3.set_xlabel("Iteration")
    ax3.set_ylabel("Mean MAE")
    ax3.set_title("Mean Regression MAE per Evaluated Rubric")
    ax3.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(save, dpi=150)
    plt.close()
    print(f"Saved {save}")

def plot_rubric_composition(results, save="/app/v2_rubric_composition.png"):
    rubric = results.get("best_rubric", {})
    items = rubric.get("items", [])
    if not items:
        return

    fig, ax = plt.subplots(figsize=(10, 6))
    labels = [i["id"] for i in items]
    n_features = [len(i.get("feature_focus", [])) for i in items]
    n_ops = [len(i.get("preferred_ops", [])) for i in items]

    x = np.arange(len(labels))
    width = 0.35
    ax.bar(x - width/2, n_features, width, label='# Features', color='steelblue')
    ax.bar(x + width/2, n_ops, width, label='# Operators', color='coral')
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=45, ha='right')
    ax.set_ylabel("Count")
    ax.set_title("Best Rubric: Constraint Composition")
    ax.legend()
    ax.grid(True, alpha=0.3, axis='y')

    plt.tight_layout()
    plt.savefig(save, dpi=150)
    plt.close()
    print(f"Saved {save}")

def plot_factor_ic_distribution(results, save="/app/v2_ic_distribution.png"):
    """Plot IC distribution of all evaluated factors."""
    hist = results.get("search_history", [])
    all_ics = []
    for h in hist:
        # We don't have per-factor data in history; get from top_factor approximation
        pass

    # Use the fact that mean_ic gives us approximate distribution
    ics = [h.get("mean_ic", 0) for h in hist if h.get("mean_ic", 0) > 0]
    if len(ics) < 2:
        return

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.hist(ics, bins=20, color='steelblue', edgecolor='black', alpha=0.7)
    ax.axvline(np.mean(ics), color='red', linestyle='--', linewidth=2, label=f'Mean={np.mean(ics):.4f}')
    ax.set_xlabel("Mean IC of Rubric-Generated Factors")
    ax.set_ylabel("Frequency")
    ax.set_title("Distribution of Factor Quality Across Rubric Evaluations")
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(save, dpi=150)
    plt.close()
    print(f"Saved {save}")

def create_report(results, save="/app/v2_report.txt"):
    lines = []
    lines.append("=" * 70)
    lines.append("AFRES v2 – Rubric-Discovery Report")
    lines.append("=" * 70)
    lines.append("")

    rubric = results.get("best_rubric", {})
    items = rubric.get("items", [])
    lines.append(f"BEST DISCOVERED RUBRIC  ({len(items)} item(s))")
    lines.append("-" * 50)
    for item in items:
        lines.append(f"  [{item['id']}] {item['description']}")
        lines.append(f"    features: {item.get('feature_focus', [])}")
        lines.append(f"    operators: {item.get('preferred_ops', [])}")
        lines.append(f"    horizon: {item.get('time_horizon_hint', '?')}  "
                     f"complexity: {item.get('complexity_hint', '?')}")
        lines.append("")

    lines.append("-" * 50)
    lines.append(f"Best fitness (mean IC): {results.get('best_fitness', 0):.4f}")
    lines.append(f"Total factors generated: {results.get('total_factors_generated', 0)}")
    lines.append(f"Valid factors: {results.get('valid_factors', 0)}")
    lines.append(f"Mean IC (all valid): {results.get('mean_ic', 0):.4f}")
    lines.append(f"Mean MAE (all valid): {results.get('mean_mae', 0):.4f}")
    lines.append(f"Mean Sharpe (all valid): {results.get('mean_sharpe', 0):.2f}")
    lines.append("")

    top = results.get("top_factor", {})
    lines.append("TOP FACTOR")
    lines.append("-" * 50)
    lines.append(f"  ID:        {top.get('id', '?')}")
    lines.append(f"  Expression: {top.get('expression', '?')}")
    lines.append(f"  IC:        {top.get('ic', 0):+.4f}")
    lines.append(f"  MAE:       {top.get('mae', 0):.4f}")
    lines.append(f"  R²:        {top.get('r2', 0):.4f}")
    lines.append(f"  Sharpe:    {top.get('sharpe', 0):.2f}")
    lines.append(f"  Returns:   {top.get('returns', 0):.2f}")
    lines.append("")

    lines.append("=" * 70)
    report = "\n".join(lines)
    with open(save, "w") as f:
        f.write(report)
    print(f"Saved {save}")
    return report

def main():
    results = load()
    plot_search_history(results)
    plot_rubric_composition(results)
    plot_factor_ic_distribution(results)
    report = create_report(results)
    print("\n" + report)

if __name__ == "__main__":
    main()