"""
Simple script to generate a Kolam image using the AI generator.
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import sys

# Add project paths
sys.path.insert(0, str(Path.cwd()))
sys.path.insert(0, str(Path.cwd() / 'models'))

from models.gan_generator import KolamGenerator
from utils.image_utils import create_synthetic_kolam
from utils.metrics import KolamDesignMetrics

def generate_ai_kolam():
    """Generate a Kolam image using the AI generator."""
    print("🎨 Generating AI Kolam Design...")
    
    # Create the generator
    generator = KolamGenerator(
        noise_dim=100,
        feature_dim=128, 
        output_channels=1,
        image_size=64
    )
    generator.eval()
    
    # Generate random noise
    noise = torch.randn(1, 100)
    
    # Generate the Kolam
    with torch.no_grad():
        generated_kolam = generator(noise)
    
    # Convert to numpy and normalize
    kolam_image = generated_kolam.squeeze().cpu().numpy()
    kolam_image = (kolam_image + 1) / 2  # Convert from [-1, 1] to [0, 1]
    kolam_image = np.clip(kolam_image, 0, 1)
    
    return kolam_image

def generate_synthetic_kolam():
    """Generate a synthetic Kolam using traditional patterns."""
    print("🎨 Generating Synthetic Kolam Design...")
    
    # Create a synthetic Kolam with medium complexity
    kolam_image = create_synthetic_kolam(size=(64, 64), complexity='medium')
    
    return kolam_image

def analyze_kolam_quality(kolam_image):
    """Analyze the quality of the generated Kolam."""
    print("📏 Analyzing Kolam Quality...")
    
    metrics = KolamDesignMetrics()
    quality = metrics.calculate_overall_quality(kolam_image)
    
    print(f"Overall Quality Score: {quality['overall_quality']:.3f}")
    print(f"Horizontal Symmetry: {quality['horizontal']:.3f}")
    print(f"Vertical Symmetry: {quality['vertical']:.3f}")
    print(f"Complexity: {quality['complexity']:.3f}")
    print(f"Balance: {quality['balance']:.3f}")
    print(f"Rhythm: {quality['rhythm']:.3f}")
    
    return quality

def display_kolam(kolam_image, title="Generated Kolam Design"):
    """Display the Kolam image."""
    plt.figure(figsize=(8, 8))
    plt.imshow(kolam_image, cmap='gray', vmin=0, vmax=1)
    plt.title(title, fontsize=16, fontweight='bold')
    plt.axis('off')
    
    # Save the image
    output_dir = Path('data/generated')
    output_dir.mkdir(parents=True, exist_ok=True)
    
    plt.savefig(output_dir / 'generated_kolam.png', dpi=300, bbox_inches='tight')
    print(f"✅ Kolam saved to: {output_dir / 'generated_kolam.png'}")
    
    plt.show()

def main():
    """Main function to generate and display Kolam."""
    print("🚀 KOLAM AI GENERATOR - IMAGE GENERATION")
    print("=" * 50)
    
    # Generate AI Kolam
    ai_kolam = generate_ai_kolam()
    print("✅ AI Kolam generated!")
    
    # Analyze quality
    quality = analyze_kolam_quality(ai_kolam)
    
    # Display the Kolam
    display_kolam(ai_kolam, "AI Generated Kolam Design")
    
    print("\n🎉 Kolam generation complete!")
    print("\nTo generate more Kolams:")
    print("1. Run this script again: python generate_kolam_demo.py")
    print("2. Use the Jupyter notebooks for interactive generation")
    print("3. Train the models with real data for better results")

if __name__ == "__main__":
    main()