import os
import sys
import torch
import json
import random
import ast
from transformers import AutoTokenizer

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from model import RecursiveCausalLM, ModelConfig, KVCache

def generate_random_graph(num_nodes, seed):
    random.seed(seed)
    nodes = [chr(65 + i) for i in range(num_nodes)]
    adj = {node: [] for node in nodes}
    
    # Connected component base
    shuffled = list(nodes)
    random.shuffle(shuffled)
    for i in range(1, num_nodes):
        u = shuffled[random.randint(0, i-1)]
        v = shuffled[i]
        adj[u].append(v)
        adj[v].append(u)
        
    # Extra random edges
    for u in nodes:
        for v in nodes:
            if u != v and v not in adj[u] and random.random() < 0.2:
                adj[u].append(v)
                adj[v].append(u)
                
    for u in adj:
        adj[u] = sorted(list(set(adj[u])))
        
    return adj

def check_connectivity(adj, start, target):
    visited = {start}
    queue = [start]
    while queue:
        curr = queue.pop(0)
        if curr == target:
            return True
        for neighbor in adj.get(curr, []):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return False

def validate_solution(adj, start, target, generation):
    # 1. Check for impossible path assertion
    if "Answer: Path is impossible." in generation:
        has_path = check_connectivity(adj, start, target)
        # If no path exists, impossible is correct
        return not has_path, "impossible_correct" if not has_path else "impossible_incorrect"
        
    # 2. Check for path list format
    if "Answer: Path is " in generation:
        try:
            idx = generation.find("Answer: Path is ")
            path_str = generation[idx + len("Answer: Path is "):].strip()
            if "<|endoftext|>" in path_str:
                path_str = path_str.split("<|endoftext|>")[0].strip()
                
            path = ast.literal_eval(path_str)
            if not isinstance(path, list) or len(path) == 0:
                return False, "format_error"
                
            if path[0] != start:
                return False, "invalid_start"
            if path[-1] != target:
                return False, "invalid_target"
                
            for i in range(len(path) - 1):
                u, v = path[i], path[i+1]
                if v not in adj.get(u, []):
                    return False, f"invalid_edge_{u}_{v}"
                    
            if len(path) != len(set(path)):
                return False, "contains_loop"
                
            return True, "correct_path"
        except Exception as e:
            return False, f"parse_exception_{type(e).__name__}"
            
    # Neither format matched or generation got truncated
    if len(generation.strip()) == 0:
        return False, "empty_generation"
    return False, "no_answer_block"

@torch.no_grad()
def greedy_decode(model, tokenizer, prompt, max_new_tokens=512, device="cuda"):
    model.eval()
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    
    # Initialize recursive multi-step KVCache
    kv_cache = KVCache(model.config, max_batch_size=1, device=device, dtype=torch.float16)
    
    use_amp = (device.type == "cuda")
    with torch.amp.autocast(device_type="cuda", enabled=use_amp, dtype=torch.float16):
        logits, _ = model(input_ids, kv_cache=kv_cache)
        
    next_token_logits = logits[0, -1, :]
    generated_tokens = []
    
    for _ in range(max_new_tokens):
        # Strict greedy selection (temperature = 0.0)
        next_token = torch.argmax(next_token_logits, dim=-1).item()
        generated_tokens.append(next_token)
        
        if next_token == tokenizer.eos_token_id:
            break
            
        curr_input = torch.tensor([[next_token]], dtype=torch.long, device=device)
        with torch.amp.autocast(device_type="cuda", enabled=use_amp, dtype=torch.float16):
            logits, _ = model(curr_input, kv_cache=kv_cache)
            next_token_logits = logits[0, -1, :]
            
    return tokenizer.decode(generated_tokens)

def main():
    print("=======================================================================")
    print("[DIAGNOSTIC] COGNITIVE LOGICAL PASS@1 RATE EVALUATOR (GRAPH TRAVERSALS)")
    print("=======================================================================\n")
    
    checkpoint_path = sys.argv[1] if len(sys.argv) > 1 else "micro_llm_200m/uct_target_sft.pt"
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Loading Model Checkpoint: {checkpoint_path}")
    print(f"Evaluation Hardware:      {device}\n")
    
    # Instantiate Model Config
    config = ModelConfig(
        vocab_size=50272,
        d_model=768,
        n_iterations=16,
        n_heads=12,
        n_kv_heads=4,
        d_ff=2048,
        max_seq_len=512
    )
    
    model = RecursiveCausalLM(config).to(device)
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    model.load_state_dict(checkpoint["model_state_dict"], strict=False)
    model.eval()
    
    tokenizer_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tokenizer")
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    
    # Generate 100 deterministic, seed-controlled evaluation problems
    num_eval_problems = 100
    problems = []
    
    # We use stable seeds 4000 to 5000 to prevent training contamination
    for i in range(num_eval_problems):
        seed = 4000 + i
        random.seed(seed)
        num_nodes = random.randint(5, 8)
        adj = generate_random_graph(num_nodes, seed)
        nodes = list(adj.keys())
        start, target = random.sample(nodes, 2)
        problems.append((adj, start, target, seed))
        
    print(f"Generated {num_eval_problems} deterministic test graphs successfully.")
    print("-----------------------------------------------------------------------\n")
    
    correct_count = 0
    categories = {}
    
    for idx, (adj, start, target, seed) in enumerate(problems):
        prompt = f"Algorithm: DFS\nQuestion: Graph: " + ", ".join([f"{u}-[{''.join(neighbors)}]" for u, neighbors in adj.items()]) + f"\nTarget: Find path from {start} to {target}.\nTrace: <search> Start"
        
        # Greedy decode response
        completion = greedy_decode(model, tokenizer, prompt, device=device)
        
        # Validate completion
        is_correct, reason = validate_solution(adj, start, target, completion)
        
        if is_correct:
            correct_count += 1
            
        categories[reason] = categories.get(reason, 0) + 1
        
        # Print a sample of results to show live trace properties
        if idx < 5:
            print(f"--- [Sample {idx+1}/5] Seed: {seed} ---")
            print(f"Start: {start} -> Target: {target}")
            print(f"Graph Adjacency: {adj}")
            print(f"Generated Response: \"{completion.strip()}\"")
            print(f"Logical Verdict:    {is_correct} ({reason})")
            print("-" * 71)
            
    pass_rate = correct_count / num_eval_problems
    
    print("\n=======================================================================")
    print("[REPORT] FINAL COGNITIVE REASONING PASS@1 METRICS REPORT")
    print("=======================================================================")
    print(f"Total Problems Tested:       {num_eval_problems}")
    print(f"Logical Pass@1 Path Rate:    {pass_rate:.4f} ({correct_count}/{num_eval_problems})")
    print("\nExecution Failure Mode Breakdown:")
    for reason, count in sorted(categories.items(), key=lambda x: x[1], reverse=True):
        percentage = count / num_eval_problems * 100
        print(f" - {reason:<30} : {count:3d} ({percentage:5.1f}%)")
    print("=======================================================================")

if __name__ == "__main__":
    main()