import torch
import time
import os
import psutil
import sys
from pathlib import Path

# Add project root to path
sys.path.append(str(Path(__file__).parent.parent))

from src.models.student import LIPEV2Student

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def estimate_flops(model, input_patches, input_landmarks):
    """
    Very rough estimate of FLOPs for the LIPEV2Student.
    """
    flops = 0
    # 1. MiniConvEmbedder (3 Conv layers + GAP)
    # Conv1: 16 filters, 3x3 kernel, input 1x8x8 -> output 16x6x6
    # FLOPs approx: 2 * Cin * K * K * Hout * Wout * Cout
    flops += 4 * (2 * 1 * 3 * 3 * 6 * 6 * 16)
    # Conv2: 16 -> 32, 6x6 -> 4x4
    flops += 4 * (2 * 16 * 3 * 3 * 4 * 4 * 32)
    # Conv3: 32 -> 64, 4x4 -> 2x2
    flops += 4 * (2 * 32 * 3 * 3 * 2 * 2 * 64)
    
    # 2. Geo MLP
    # Linear 1: 956 -> 256
    flops += 2 * 956 * 256
    # Linear 2: 256 -> 256
    flops += 2 * 256 * 256
    
    # 3. Heads
    # Pitch: 256 -> 64 -> 90
    flops += 2 * 256 * 64 + 2 * 64 * 90
    # Yaw: 256 -> 64 -> 90
    flops += 2 * 256 * 64 + 2 * 64 * 90
    
    return flops

def profile_model():
    model = LIPEV2Student()
    model.eval()
    
    params = count_parameters(model)
    print(f"Total Trainable Parameters: {params:,}")
    
    dummy_patches = torch.randn(1, 4, 8, 8)
    dummy_landmarks = torch.randn(1, 956)
    
    # FLOPs Estimation
    flops = estimate_flops(model, dummy_patches, dummy_landmarks)
    print(f"Estimated FLOPs per frame (State A): {flops / 1e6:.2f} MFLOPs ({flops / 1e9:.4f} GFLOPs)")
    
    # RAM Measurement
    process = psutil.Process(os.getpid())
    mem_before = process.memory_info().rss / 1024 / 1024 # MB
    
    # Run some inferences
    with torch.no_grad():
        for _ in range(100):
            _ = model(dummy_patches, dummy_landmarks, state='A')
            
    mem_after = process.memory_info().rss / 1024 / 1024 # MB
    print(f"Memory Usage (RSS): {mem_after:.2f} MB (Baseline: {mem_before:.2f} MB)")
    
    # Latency
    start_time = time.time()
    num_iters = 500
    with torch.no_grad():
        for _ in range(num_iters):
            _ = model(dummy_patches, dummy_landmarks, state='A')
    end_time = time.time()
    
    avg_latency = (end_time - start_time) / num_iters * 1000
    print(f"Average Inference Latency (State A): {avg_latency:.4f} ms")
    print(f"Estimated Max FPS (State A): {1000 / avg_latency:.2f}")

if __name__ == '__main__':
    profile_model()