import io
import json
import numpy as np
import sys
import os
import re
import importlib
import random
from datasets import load_dataset
from hoho2025.metric_helper import hss

print("Loading dataset from Hugging Face streaming (usm3d/hoho22k_2026_trainval)...")
dataset = load_dataset('usm3d/hoho22k_2026_trainval', split='train', streaming=True, trust_remote_code=True)
# Fetch first 3 samples into memory to speed up search
samples = []
for idx, s in enumerate(dataset):
    if idx >= 100:
        break
    samples.append(s)

def set_hyperparameters(eps, min_len, min_support, max_angle, max_dist, min_degree, min_view):
    with open('script.py', 'r') as f:
        content = f.read()

    content = re.sub(r'VERTEX_MERGE_EPS = .+', f'VERTEX_MERGE_EPS = {eps}', content)
    content = re.sub(r'EDGE_MIN_LENGTH = .+', f'EDGE_MIN_LENGTH = {min_len}', content)
    content = re.sub(r'EDGE_MIN_SUPPORT_IMAGES = .+', f'EDGE_MIN_SUPPORT_IMAGES = {min_support}', content)
    content = re.sub(r'EDGE_MAX_ANGLE_DEG = .+', f'EDGE_MAX_ANGLE_DEG = {max_angle}', content)
    content = re.sub(r'VERTEX_MAX_COLMAP_DIST = .+', f'VERTEX_MAX_COLMAP_DIST = {max_dist}', content)
    content = re.sub(r'VERTEX_MIN_EDGE_DEGREE = .+', f'VERTEX_MIN_EDGE_DEGREE = {min_degree}', content)
    content = re.sub(r'VERTEX_MIN_VIEW_COUNT = .+', f'VERTEX_MIN_VIEW_COUNT = {min_view}', content)

    with open('script.py', 'w') as f:
        f.write(content)

import script

def evaluate():
    importlib.reload(script)
    scores = []
    for sample in samples:
        try:
            pred_v, pred_e, _ = script.predict_wireframe_safely(sample)
        except Exception:
            pred_v, pred_e = np.zeros((2, 3)), [(0, 1)]
            
        gt_v = sample.get('wf_vertices')
        gt_e = sample.get('wf_edges')
        res = hss(pred_v, pred_e, gt_v, gt_e)
        scores.append(res.hss)
    
    # Cache cleaning to prevent memory leaks over many trials
    import gc
    gc.collect()
    
    return sum(scores) / len(scores) if scores else 0

best_score = -1.0
best_params = None

# Baseline is around 0.34
space = {
    'eps': [0.1, 0.5, 1.0],
    'min_len': [0.0, 0.2, 0.5],
    'min_support': [1, 2],
    'max_angle': [25.0, 45.0],
    'max_dist': [2.0, 4.0],
    'min_degree': [1, 2],
    'min_view': [1, 2]
}

trials = 50
print(f"Running Random Search for {trials} trials...")

for i in range(trials):
    p = {k: random.choice(v) for k, v in space.items()}
    set_hyperparameters(
        p['eps'], p['min_len'], p['min_support'], 
        p['max_angle'], p['max_dist'], p['min_degree'], p['min_view']
    )
    score = evaluate()
    print(f"Trial {i+1} Score: {score:.4f} | Params: {p}")
    if score > best_score:
        best_score = score
        best_params = p

print(f"\n--- BEST RESULT ---")
print(f"Best HSS Score: {best_score:.4f}")
print("Best Parameters:")
for k, v in best_params.items():
    print(f"  {k}: {v}")

# Restore best params to script.py
set_hyperparameters(
    best_params['eps'], best_params['min_len'], best_params['min_support'], 
    best_params['max_angle'], best_params['max_dist'], best_params['min_degree'], best_params['min_view']
)