"""
Complete Analysis and Test Framework for Tong Nian's Challenge

This script analyzes the challenge requirements and tests potential solutions.
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

import hashlib
import io
import zipfile
from pathlib import Path
from PIL import Image
import numpy as np
import torch

# Import the actual verification and model code
from config import load_config
from model import get_model_bundle, predict_top1, set_deterministic
from verification import Verifier, md5_hex, load_gray_from_bytes


def analyze_references():
    """Analyze the reference images and requirements."""
    config = load_config()
    refs_dir = Path("src/data/refs")
    
    print("="*80)
    print("CHALLENGE REQUIREMENTS ANALYSIS")
    print("="*80)
    print(f"\nTarget classes: {config.target_class_ids}")
    print(f"Reference classes: {config.reference_class_ids}")
    print(f"Pixel budgets: {config.budgets}")
    print(f"Flag: {config.flag}")
    
    print("\n" + "="*80)
    print("REFERENCE IMAGES ANALYSIS")
    print("="*80)
    
    # Set up model for prediction
    set_deterministic(0)
    bundle = get_model_bundle(device="cpu")
    
    for i in range(10):
        ref_path = refs_dir / f"ref_{i:02d}.png"
        if not ref_path.exists():
            print(f"\nPair {i:02d}: MISSING")
            continue
            
        img = Image.open(ref_path).convert("L")
        arr = np.array(img)
        
        # Get model prediction
        pred = predict_top1(img, bundle=bundle)
        
        print(f"\nPair {i:02d}:")
        print(f"  Shape: {arr.shape}")
        print(f"  Pixel range: [{arr.min()}, {arr.max()}]")
        print(f"  Current classification: {pred['label']} (prob: {pred['prob']:.4f})")
        print(f"  Required img1 class: {config.reference_class_ids[i]}")
        print(f"  Required img2 class: {config.target_class_ids[i]}")
        print(f"  Pixel budget for img2: {config.budgets[i]} pixels")
        print(f"  Challenge: img1==img2 (same bytes) but different classification!")


def test_identical_classification():
    """
    Test if calling predict_top1 twice on the same image gives the same result.
    This tests for any non-determinism in the model.
    """
    print("\n" + "="*80)
    print("NON-DETERMINISM TEST")
    print("="*80)
    
    set_deterministic(0)
    bundle = get_model_bundle(device="cpu")
    
    refs_dir = Path("src/data/refs")
    ref_path = refs_dir / "ref_00.png"
    
    # Load image once
    with open(ref_path, 'rb') as f:
        img_bytes = f.read()
    
    # Test multiple predictions
    print("\nTesting if same bytes give same prediction...")
    results = []
    for trial in range(5):
        img = load_gray_from_bytes(img_bytes)
        pred = predict_top1(img, bundle=bundle)
        results.append(pred['id'])
        print(f"  Trial {trial+1}: class={pred['id']}, prob={pred['prob']:.6f}")
    
    if len(set(results)) == 1:
        print("\n✓ Model is deterministic - all predictions identical")
    else:
        print("\n✗ Model shows non-determinism!")
        print(f"  Got different results: {set(results)}")


def test_context_manager_behavior():
    """
    Test the exact pattern used in verification.py to see if there's any quirk.
    """
    print("\n" + "="*80)
    print("CONTEXT MANAGER PATTERN TEST")
    print("="*80)
    
    refs_dir = Path("src/data/refs")
    ref_path = refs_dir / "ref_00.png"
    
    with open(ref_path, 'rb') as f:
        img_bytes = f.read()
    
    print(f"\nOriginal file MD5: {md5_hex(img_bytes)}")
    
    # Load twice using the same pattern as verification
    img1 = load_gray_from_bytes(img_bytes)
    img2 = load_gray_from_bytes(img_bytes)
    
    arr1 = np.array(img1)
    arr2 = np.array(img2)
    
    print(f"img1 shape: {arr1.shape}, dtype: {arr1.dtype}")
    print(f"img2 shape: {arr2.shape}, dtype: {arr2.dtype}")
    print(f"Arrays equal: {np.array_equal(arr1, arr2)}")
    print(f"Pixel differences: {np.sum(arr1 != arr2)}")
    
    # Test if PIL Image objects are somehow different
    print(f"\nPIL Image comparison:")
    print(f"  Same object: {img1 is img2}")
    print(f"  img1.mode: {img1.mode}, img2.mode: {img2.mode}")
    print(f"  img1.size: {img1.size}, img2.size: {img2.size}")


def explore_solution_approaches():
    """
    Document potential approaches to solve the challenge.
    """
    print("\n" + "="*80)
    print("POTENTIAL SOLUTION APPROACHES")
    print("="*80)
    
    approaches = [
        ("1. MD5 Collision Attack",
         "Create two different PNG files with the same MD5 hash using cryptographic collision attacks.",
         "Very difficult; would require advanced crypto techniques like chosen-prefix collision."),
        
        ("2. PNG Format Tricks",
         "Exploit PNG chunk parsing differences where same bytes decode differently.",
         "PNG format is deterministic; same bytes should always give same image."),
        
        ("3. ZIP File Tricks", 
         "Use ZIP format features like directory entries pointing to same data.",
         "But verification reads files separately, so this won't help."),
        
        ("4. Model Non-Determinism",
         "Find non-deterministic behavior in PyTorch model or preprocessing.",
         "Model uses eval() mode and set_deterministic(), but might have edge cases."),
        
        ("5. PIL Lazy Loading",
         "Exploit PIL's lazy loading where Image.convert() might have state.",
         "Unlikely - PIL is generally deterministic."),
        
        ("6. Preprocessing Non-Determinism",
         "Find floating-point differences in preprocessing that affect classification.",
         "Model preprocessing includes resize, which might have anti-aliasing differences."),
        
        ("7. The Name is a Clue",
         "'Nothing Ever Changes' - maybe the solution is that LITERALLY nothing changes?",
         "Perhaps img1 and img2 are meant to be identical files, and the challenge is impossible?"),
        
        ("8. There's a Bug in Verification",
         "Maybe there's a logical error in verification.py that makes it passable.",
         "Worth reviewing the code line by line for any oversights."),
    ]
    
    for title, description, notes in approaches:
        print(f"\n{title}")
        print(f"  {description}")
        print(f"  → {notes}")


def test_verification_locally():
    """
    Test the verification logic with dummy data.
    """
    print("\n" + "="*80)
    print("LOCAL VERIFICATION TEST")
    print("="*80)
    
    print("\nAttempting to create a simple test case...")
    print("Creating a ZIP with identical files for pair_00...")
    
    # Create a simple test
    refs_dir = Path("src/data/refs")
    ref_path = refs_dir / "ref_00.png"
    
    with open(ref_path, 'rb') as f:
        ref_bytes = f.read()
    
    # Create a ZIP with img1 and img2 being the same file
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zf:
        for i in range(10):
            # Use the same reference file for both img1 and img2
            ref_i_path = refs_dir / f"ref_{i:02d}.png"
            if ref_i_path.exists():
                with open(ref_i_path, 'rb') as f:
                    data = f.read()
                zf.writestr(f"pair_{i:02d}_img1.png", data)
                zf.writestr(f"pair_{i:02d}_img2.png", data)
    
    zip_bytes = zip_buffer.getvalue()
    
    # Test verification
    config = load_config()
    set_deterministic(0)
    bundle = get_model_bundle(device="cpu")
    
    def predict_fn(image):
        return predict_top1(image, bundle=bundle)
    
    verifier = Verifier(config=config, predict_fn=predict_fn)
    
    print("\nRunning verification on test ZIP...")
    result = verifier.verify_zip(zip_bytes)
    
    print(f"\nVerification result: {result}")
    if not result:
        print("As expected, verification failed because:")
        print("  - img1 and img2 have same bytes (MD5 matches) ✓")
        print("  - img1 and img2 have same pixels ✓")
        print("  - But img1 and img2 classify the same way ✗")
        print("  - Need img2 to classify differently than img1!")


if __name__ == "__main__":
    try:
        analyze_references()
        test_identical_classification()
        test_context_manager_behavior()
        explore_solution_approaches()
        test_verification_locally()
        
        print("\n" + "="*80)
        print("SUMMARY")
        print("="*80)
        print("""
Tong Nian's claim: "Create adversarial examples without changing anything"

The Challenge:
- Submit pairs of images (img1, img2) where:
  * img1_bytes == img2_bytes (same MD5 hash - literally identical bytes!)
  * img1 must classify as reference class
  * img2 must classify as target class (different from reference)
  * Both must be valid PNG images

The Paradox:
- How can identical bytes produce different classifications?
- This seems impossible unless there's:
  a) Non-determinism in the model/preprocessing
  b) A bug in the verification logic
  c) Some extremely clever PNG/ZIP trick we haven't thought of
  d) The challenge is actually unsolvable (philosophical point about ML)

To proceed, we would need to:
1. Carefully analyze the model for any non-deterministic behavior
2. Study the preprocessing pipeline for floating-point quirks
3. Investigate advanced PNG format features
4. Consider if the challenge itself is a commentary on reproducibility in ML

The application is running at: http://localhost:5001
Try submitting test cases to understand the exact failure modes!
        """)
        
    except Exception as e:
        print(f"\nError during analysis: {e}")
        import traceback
        traceback.print_exc()