import os
import time

import anthropic
from metaprompt import RESEARCH_PRINCIPLES, MODELS, NUM_IDEAS, LOG_FILE, GENERATION_DIRECTORY_PATH, IMPORT_STRUCTURE_PREFIX, SUMMARIZE_IMMEDIATELY
from generate import AlgoGen
from describe import ModelAnalyzer

ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
anthopic_client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)
algo_gen = AlgoGen(anthropic_client=anthopic_client, log_file=LOG_FILE)
analyzer = ModelAnalyzer(anthropic_client=anthopic_client)
summaries = os.path.join(GENERATION_DIRECTORY_PATH, "model_summaries.txt")
# benchmark_suite = BenchmarkSuite()
total_models = 0
for m in MODELS[1:]:
    init_file_path = os.path.join(GENERATION_DIRECTORY_PATH, '__init__.py')
    if not os.path.exists(init_file_path):
        Warning("__init__.py doesnt exist, modules might not be imported.")
    if(IMPORT_STRUCTURE_PREFIX == ""): 
            Warning("IMPORT_STRUCTURE_PREFIX is empty, imports may fail.")

    print(f"Generating ideas for model {m}...")
    ideas_raw = algo_gen.gen(f"""Principles of Machine Learning: {RESEARCH_PRINCIPLES}
    Inspired by these principles, come up with 10 very creative ideas for how to vary the {m} classifier for better performance.
    
    IMPORTANT: Format your response as exactly one idea per line, with each line starting with "IDEA: " followed by the idea description.
    Example format:
    IDEA: Use adaptive compression based on local density
    IDEA: Implement multi-level abstraction layers
    IDEA: Apply entropy-guided feature selection
    
    Do not include any other text, just the 10 IDEA lines.""")
    
    ideas_list = []
    for line in ideas_raw.split('\n'):
        line = line.strip()
        if line.startswith('IDEA: '):
            idea = line[6:].strip()
            if idea:
                ideas_list.append(idea)
    
    print(f"Extracted {len(ideas_list)} ideas")
    
    if not ideas_list:
        print(f"Warning: No ideas extracted for {m}")
        continue
    
    ideas_list = ideas_list[:NUM_IDEAS]    

    print(f"Generating implementations for {len(ideas_list)} ideas (limited to {NUM_IDEAS})...")
    generated_files_result = algo_gen.parallel_genML(ideas_list)
    num_successful_models = sum(x is not None for x in generated_files_result)
    if num_successful_models > 0: 
        print(f"Successfully generated {num_successful_models} models")
        if SUMMARIZE_IMMEDIATELY:
            for result in generated_files_result:
                if result is None:
                    continue
                filename, classname, idea = result
                if filename:
                    description = analyzer.describe_single(GENERATION_DIRECTORY_PATH, filename)
                    with open(summaries, "a", encoding="utf-8") as f:
                        f.write(description + "\n")

    total_models += num_successful_models
    print(f"Results: {generated_files_result}")
print(f"Total models generated: {total_models}")
