"""Optimized weight finder using random search on cached data."""
import json
import random
import math
from pathlib import Path
from collections import defaultdict
import itertools

import sys
sys.path.insert(0, '/home/brettanthonysjoberg179/trifecta-bro-hf-space')

from trifecta_bro.data.models import RaceModel, RunnerModel
from trifecta_bro.model.scoring import score_runner
from trifecta_bro.model.pace_analysis import classify_pace
from trifecta_bro.model.distance_profiles import classify_distance

DATA_DIR = Path('/home/brettanthonysjoberg179/trifecta-bro-hf-space/data')
CACHE_DIR = DATA_DIR / 'cache'
RESULTS_DIR = DATA_DIR / 'results'

random.seed(42)


def load_cached_races(date: str) -> list[dict]:
    """Load all cached form data for a date."""
    races = []
    for f in CACHE_DIR.glob('*.json'):
        try:
            with open(f) as fp:
                data = json.load(fp)
            if data.get('date') == date:
                races.append(data)
        except:
            continue
    return races


def load_results(date: str) -> dict:
    """Load actual results for a date."""
    result_lookup = {}
    result_path = RESULTS_DIR / f"{date}-ra.json"
    if not result_path.exists():
        result_path = RESULTS_DIR / f"{date}.json"
    if not result_path.exists():
        return result_lookup
    
    with open(result_path) as f:
        rdata = json.load(f)
    
    if isinstance(rdata, dict):
        if 'tracks' in rdata:
            for track, races in rdata['tracks'].items():
                if isinstance(races, dict):
                    for rn, r in races.items():
                        runners = r.get('runners', [])
                        winners = []
                        for runner in runners:
                            pos = runner.get('position')
                            if pos and pos <= 3:
                                winners.append((pos, runner['number']))
                        winners.sort()
                        actual = [w[1] for w in winners]
                        if actual:
                            result_lookup[(track, int(rn))] = actual
                elif isinstance(races, list):
                    for r in races:
                        if 'trifecta' in r:
                            result_lookup[(track, r.get('race', 0))] = [int(x) for x in r['trifecta']]
        else:
            for track, races in rdata.items():
                if isinstance(races, list):
                    for r in races:
                        if 'trifecta' in r:
                            result_lookup[(track, r['race'])] = [int(x) for x in r['trifecta']]
    
    return result_lookup


def score_race_with_weights(race_data: dict, weights: dict) -> list[int] | None:
    """Score a race with given weights, return trifecta."""
    try:
        runners = []
        for r in race_data.get('runners', []):
            if r.get('scratched'):
                continue
            runner = RunnerModel(
                number=r['number'],
                name=r.get('name', ''),
                jockey=r.get('jockey'),
                trainer=r.get('trainer'),
                weight=r.get('weight'),
                barrier=r.get('barrier'),
                age=r.get('age'),
                sex=r.get('sex'),
                form=r.get('form', ''),
                last20_starts=r.get('last20Starts', ''),
                stats=r.get('stats', {}),
            )
            runners.append(runner)
        
        if len(runners) < 3:
            return None
        
        race = RaceModel(
            date=race_data.get('date', ''),
            track=race_data.get('track', ''),
            track_slug=race_data.get('slug', race_data.get('track', '').lower()),
            race_number=race_data.get('raceNumber', 0),
            race_name=race_data.get('raceName', ''),
            distance=race_data.get('distance'),
            condition=race_data.get('condition'),
            race_class=race_data.get('raceClass'),
            abandoned=race_data.get('abandoned', False),
            start_time=race_data.get('startTime'),
            prize_money=str(race_data.get('prizeMoney', '')),
            number_of_runners=race_data.get('numberOfRunners', len(runners)),
            runners=runners,
        )
        
        pace = classify_pace(race, runners)
        
        scored = []
        for r in runners:
            sc = score_runner(r, race, pace, weights)
            scored.append((r.number, sc['score']))
        
        scored.sort(key=lambda x: x[1], reverse=True)
        return [s[0] for s in scored[:3]]
    except:
        return None


def evaluate_weights_on_date(date: str, weights: dict, result_lookup: dict) -> dict:
    """Evaluate weights on a single date."""
    races = load_cached_races(date)
    stats = {"total": 0, "top1": 0, "exact": 0, "box": 0}
    
    for race_data in races:
        track = race_data.get('track', '')
        race_num = race_data.get('raceNumber', 0)
        actual = result_lookup.get((track, race_num))
        
        if not actual:
            continue
        
        pred = score_race_with_weights(race_data, weights)
        if not pred:
            continue
        
        stats["total"] += 1
        if pred == actual:
            stats["exact"] += 1
        if set(pred) == set(actual):
            stats["box"] += 1
        if pred[0] == actual[0]:
            stats["top1"] += 1
    
    return stats


def evaluate_weights_all_dates(dates: list[str], weights: dict) -> dict:
    """Evaluate weights across all dates."""
    agg = {"total": 0, "top1": 0, "exact": 0, "box": 0}
    
    for date in dates:
        result_lookup = load_results(date)
        stats = evaluate_weights_on_date(date, weights, result_lookup)
        agg["total"] += stats["total"]
        agg["top1"] += stats["top1"]
        agg["exact"] += stats["exact"]
        agg["box"] += stats["box"]
    
    agg["top1_pct"] = agg["top1"] / agg["total"] * 100 if agg["total"] > 0 else 0
    return agg


def random_weights() -> dict:
    """Generate random weights that sum to ~1.0."""
    weights = {}
    for k in ["form", "class", "distance", "track", "track_distance", 
              "condition", "jockey", "fitness", "barrier", "weight", "pace"]:
        weights[k] = random.random()
    
    total = sum(weights.values())
    return {k: v / total for k, v in weights.items()}


def random_search(dates: list[str], n_iterations: int = 200) -> list[dict]:
    """Random search for optimal weights."""
    results = []
    
    # Also test the base weights
    base_weights = {
        "form": 0.18, "class": 0.14, "distance": 0.10, "track": 0.10,
        "track_distance": 0.10, "condition": 0.05, "jockey": 0.06,
        "fitness": 0.08, "barrier": 0.08, "weight": 0.08, "pace": 0.03,
    }
    
    # Evaluate base
    base_stats = evaluate_weights_all_dates(dates, base_weights)
    results.append({"weights": base_weights, "stats": base_stats, "top1_pct": base_stats["top1_pct"]})
    print(f"Base weights: {base_stats['top1_pct']:.1f}%")
    
    # Random search
    for i in range(n_iterations):
        weights = random_weights()
        stats = evaluate_weights_all_dates(dates, weights)
        
        results.append({"weights": weights, "stats": stats, "top1_pct": stats["top1_pct"]})
        
        if (i + 1) % 50 == 0:
            best_so_far = max(results, key=lambda x: x["top1_pct"])
            print(f"  Iteration {i+1}: best = {best_so_far['top1_pct']:.1f}%")
    
    results.sort(key=lambda x: x["top1_pct"], reverse=True)
    return results


def coordinate_descent(dates: list[str], base_weights: dict, n_rounds: int = 10) -> dict:
    """Optimize one parameter at a time while holding others fixed."""
    weights = dict(base_weights)
    
    param_ranges = {
        "form": [0.10, 0.15, 0.20, 0.25, 0.30],
        "class": [0.08, 0.12, 0.16, 0.20],
        "distance": [0.04, 0.08, 0.12, 0.16],
        "track": [0.04, 0.08, 0.12, 0.16],
        "track_distance": [0.04, 0.08, 0.12, 0.16],
        "condition": [0.02, 0.05, 0.08],
        "jockey": [0.02, 0.05, 0.08, 0.10],
        "fitness": [0.04, 0.08, 0.12],
        "barrier": [0.04, 0.08, 0.12, 0.16],
        "weight": [0.04, 0.08, 0.12],
        "pace": [0.02, 0.04, 0.06, 0.08],
    }
    
    current_best = evaluate_weights_all_dates(dates, weights)
    current_pct = current_best["top1_pct"]
    print(f"Starting coordinate descent from {current_pct:.1f}%")
    
    for round_num in range(n_rounds):
        improved = False
        
        for param in param_ranges:
            best_value = weights[param]
            best_pct = current_pct
            
            for value in param_ranges[param]:
                test_weights = dict(weights)
                test_weights[param] = value
                
                # Normalize
                total = sum(test_weights.values())
                test_weights = {k: v / total for k, v in test_weights.items()}
                
                stats = evaluate_weights_all_dates(dates, test_weights)
                
                if stats["top1_pct"] > best_pct:
                    best_pct = stats["top1_pct"]
                    best_value = value
                    improved = True
            
            weights[param] = best_value
        
        # Normalize after each round
        total = sum(weights.values())
        weights = {k: v / total for k, v in weights.items()}
        
        current_best = evaluate_weights_all_dates(dates, weights)
        current_pct = current_best["top1_pct"]
        
        print(f"  Round {round_num + 1}: {current_pct:.1f}%")
        
        if not improved:
            print(f"  No improvement, stopping")
            break
    
    return {"weights": weights, "stats": current_best, "top1_pct": current_pct}


def main():
    dates = ["2026-08-07", "2026-08-08", "2026-08-09", "2026-08-14"]
    
    print("=" * 60)
    print("WEIGHT OPTIMIZATION")
    print("=" * 60)
    
    # Test base weights first
    base_weights = {
        "form": 0.18, "class": 0.14, "distance": 0.10, "track": 0.10,
        "track_distance": 0.10, "condition": 0.05, "jockey": 0.06,
        "fitness": 0.08, "barrier": 0.08, "weight": 0.08, "pace": 0.03,
    }
    
    print("\n--- Base weights ---")
    base_stats = evaluate_weights_all_dates(dates, base_weights)
    print(f"Top 1: {base_stats['top1_pct']:.1f}% ({base_stats['top1']}/{base_stats['total']})")
    
    # Random search
    print("\n--- Random search ---")
    random_results = random_search(dates, n_iterations=200)
    
    print("\n--- Top 5 from random search ---")
    for i, r in enumerate(random_results[:5]):
        w = r['weights']
        print(f"{i+1}. Top1: {r['top1_pct']:.1f}% | "
              f"form={w['form']:.2f} class={w['class']:.2f} dist={w['distance']:.2f} "
              f"track={w['track']:.2f} td={w['track_distance']:.2f} "
              f"fitness={w['fitness']:.2f} barrier={w['barrier']:.2f} "
              f"weight={w['weight']:.2f} pace={w['pace']:.2f}")
    
    # Coordinate descent from best random
    best_random = random_results[0]
    print("\n--- Coordinate descent from best random ---")
    cd_result = coordinate_descent(dates, best_random["weights"])
    
    print("\n--- Final optimized weights ---")
    print(f"Top 1: {cd_result['top1_pct']:.1f}% ({cd_result['stats']['top1']}/{cd_result['stats']['total']})")
    
    w = cd_result['weights']
    print(f"form={w['form']:.2f}, class={w['class']:.2f}, distance={w['distance']:.2f}, "
          f"track={w['track']:.2f}, track_distance={w['track_distance']:.2f}, "
          f"condition={w['condition']:.2f}, jockey={w['jockey']:.2f}, "
          f"fitness={w['fitness']:.2f}, barrier={w['barrier']:.2f}, "
          f"weight={w['weight']:.2f}, pace={w['pace']:.2f}")
    
    # Per-date breakdown
    print("\nPer-date breakdown:")
    for date in dates:
        result_lookup = load_results(date)
        stats = evaluate_weights_on_date(date, cd_result["weights"], result_lookup)
        if stats['total'] > 0:
            print(f"  {date}: {stats['total']} races, {stats['top1']} top1 ({stats['top1']/stats['total']*100:.0f}%)")
    
    return cd_result


if __name__ == "__main__":
    main()