import numpy as np
import pickle
import os
import sys
import time

# ==================== PATHS ====================
MODEL_FOLDER = "/storage/emulated/0/AI_Photos/models/"
MODEL_FILE = os.path.join(MODEL_FOLDER, "model.pkl")

print("="*60)
print("📊 NEURAL NETWORK PARAMETER MEASURER")
print("="*60)

# ==================== LOAD MODEL ====================
def load_model():
    if not os.path.exists(MODEL_FILE):
        print(f"\n❌ File not found: {MODEL_FILE}")
        return None
    
    try:
        print(f"\n📂 Loading: {os.path.basename(MODEL_FILE)}")
        with open(MODEL_FILE, 'rb') as f:
            data = pickle.load(f)
        print("✅ Model loaded!")
        return data
    except Exception as e:
        print(f"❌ Error: {e}")
        return None

# ==================== MEASUREMENTS ====================
def measure_model(data):
    print("\n" + "="*60)
    print("📊 MODEL PARAMETERS")
    print("="*60)
    
    # 1. Basic parameters
    input_size = data.get('input_size', 0)
    hidden_size = data.get('hidden_size', 0)
    output_size = data.get('output_size', 0)
    image_size = data.get('image_size', 40)
    
    print(f"\n🏗️ ARCHITECTURE:")
    print(f"   📐 Input layer:  {input_size} neurons")
    print(f"   📐 Hidden layer:  {hidden_size} neurons")
    print(f"   📐 Output layer: {output_size} neurons")
    print(f"   🖼️ Image size: {image_size}x{image_size} pixels")
    
    # 2. Weights
    W1 = data.get('W1')
    b1 = data.get('b1')
    W2 = data.get('W2')
    b2 = data.get('b2')
    
    if W1 is not None:
        # Number of parameters
        params_W1 = W1.size
        params_b1 = b1.size if b1 is not None else 0
        params_W2 = W2.size
        params_b2 = b2.size if b2 is not None else 0
        total_params = params_W1 + params_b1 + params_W2 + params_b2
        
        print(f"\n📊 PARAMETER COUNT:")
        print(f"   W1 (input→hidden):  {params_W1:,} parameters")
        print(f"   b1 (bias):          {params_b1:,} parameters")
        print(f"   W2 (hidden→output): {params_W2:,} parameters")
        print(f"   b2 (bias):          {params_b2:,} parameters")
        print(f"   📦 TOTAL:           {total_params:,} parameters")
        
        # 3. Weight shapes
        print(f"\n📏 WEIGHT SHAPES:")
        print(f"   W1: {W1.shape}")
        print(f"   b1: {b1.shape if b1 is not None else 'None'}")
        print(f"   W2: {W2.shape}")
        print(f"   b2: {b2.shape if b2 is not None else 'None'}")
    
    # 4. File size
    file_size = os.path.getsize(MODEL_FILE)
    print(f"\n💾 FILE SIZE:")
    print(f"   {file_size / 1024:.2f} KB")
    print(f"   {file_size / (1024*1024):.2f} MB")
    print(f"   {file_size} bytes")
    
    # 5. Data types
    if W1 is not None:
        print(f"\n🔢 DATA TYPES:")
        print(f"   W1: {W1.dtype}")
        print(f"   W2: {W2.dtype}")
        if b1 is not None:
            print(f"   b1: {b1.dtype}")
        if b2 is not None:
            print(f"   b2: {b2.dtype}")
    
    # 6. Weight statistics
    if W1 is not None:
        print(f"\n📊 WEIGHT STATISTICS:")
        print(f"   W1 - min: {W1.min():.6f}, max: {W1.max():.6f}, mean: {W1.mean():.6f}")
        print(f"   W2 - min: {W2.min():.6f}, max: {W2.max():.6f}, mean: {W2.mean():.6f}")
        if b1 is not None:
            print(f"   b1 - min: {b1.min():.6f}, max: {b1.max():.6f}, mean: {b1.mean():.6f}")
        if b2 is not None:
            print(f"   b2 - min: {b2.min():.6f}, max: {b2.max():.6f}, mean: {b2.mean():.6f}")
    
    # 7. Memory
    if W1 is not None:
        memory = (W1.nbytes + W2.nbytes + 
                 (b1.nbytes if b1 is not None else 0) + 
                 (b2.nbytes if b2 is not None else 0))
        print(f"\n🧠 WEIGHT MEMORY:")
        print(f"   Total: {memory / 1024:.2f} KB")
        print(f"   Total: {memory / (1024*1024):.3f} MB")
    
    # 8. Generation speed
    print(f"\n⏱️ GENERATION TEST:")
    try:
        # Create class for testing
        from generate_only import MyNeuralNetwork
        nn = MyNeuralNetwork(input_size, hidden_size, output_size)
        nn.W1 = W1
        nn.b1 = b1
        nn.W2 = W2
        nn.b2 = b2
        
        # Test speed
        times = []
        for _ in range(10):
            noise = np.random.randn(1, input_size) * 2.5
            start = time.time()
            output = nn.forward(noise)
            end = time.time()
            times.append(end - start)
        
        avg_time = np.mean(times) * 1000
        print(f"   ⚡ Average generation time: {avg_time:.2f} ms")
        print(f"   ⚡ Fastest: {np.min(times) * 1000:.2f} ms")
        print(f"   ⚡ Slowest: {np.max(times) * 1000:.2f} ms")
        print(f"   🎨 Can generate ~{int(1000/avg_time)} images/sec")
    except:
        print("   ❌ Could not test speed")
    
    # 9. FINAL RESULT
    print("\n" + "="*60)
    print("📋 FINAL CHARACTERISTICS:")
    print("="*60)
    
    # Generate name string
    if W1 is not None:
        name_parts = [
            f"AI_{image_size}x{image_size}",
            f"Params_{total_params:,}",
            f"Hidden_{hidden_size}",
            f"Size_{file_size/(1024*1024):.1f}MB"
        ]
        model_name = "_".join(name_parts)
        print(f"\n📛 MODEL NAME:")
        print(f"   {model_name}")
        print(f"\n   Use this for filename:")
        print(f"   📄 {model_name}.pkl")
    
    return model_name if W1 is not None else None

# ==================== SAVE REPORT ====================
def save_report(data, model_name):
    if model_name is None:
        return
    
    report = f"""================================
MODEL REPORT
================================

📛 NAME: {model_name}

🏗️ ARCHITECTURE:
   Input layer:  {data.get('input_size', 0)} neurons
   Hidden layer:  {data.get('hidden_size', 0)} neurons
   Output layer: {data.get('output_size', 0)} neurons
   Image size: {data.get('image_size', 40)}x{data.get('image_size', 40)} pixels

📊 PARAMETERS:
   W1: {data['W1'].size:,} parameters
   b1: {data['b1'].size:,} parameters
   W2: {data['W2'].size:,} parameters
   b2: {data['b2'].size:,} parameters
   TOTAL: {data['W1'].size + data['b1'].size + data['W2'].size + data['b2'].size:,} parameters

💾 FILE SIZE: {os.path.getsize(MODEL_FILE) / (1024*1024):.2f} MB

📏 WEIGHT SHAPES:
   W1: {data['W1'].shape}
   b1: {data['b1'].shape}
   W2: {data['W2'].shape}
   b2: {data['b2'].shape}

================================
Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}
================================
"""
    
    # Save to file
    report_path = os.path.join(MODEL_FOLDER, "model_report.txt")
    with open(report_path, 'w') as f:
        f.write(report)
    
    print(f"\n📄 Report saved: {report_path}")
    
    # Save model name to file
    name_path = os.path.join(MODEL_FOLDER, "model_name.txt")
    with open(name_path, 'w') as f:
        f.write(model_name)
    print(f"📛 Model name saved: {name_path}")

# ==================== MAIN ====================
def main():
    # 1. Load model
    data = load_model()
    if data is None:
        return
    
    # 2. Measure parameters
    model_name = measure_model(data)
    
    # 3. Save report
    if model_name:
        save_report(data, model_name)
    
    print("\n" + "="*60)
    print("✅ MEASUREMENT COMPLETE!")
    print("="*60)

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n❌ Error: {e}")