import gradio as gr
import pandas as pd
import numpy as np
import io
import base64
from typing import Optional, Tuple
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import warnings
warnings.filterwarnings("ignore")

# Import Mostly AI SDK
try:
    from mostlyai.sdk import MostlyAI
    MOSTLY_AI_AVAILABLE = True
except ImportError:
    MOSTLY_AI_AVAILABLE = False
    print("Warning: Mostly AI SDK not available. Please install with: pip install mostlyai[local]")

class SyntheticDataGenerator:
    def __init__(self):
        self.mostly = None
        self.generator = None
        self.original_data = None
        
    def initialize_mostly_ai(self):
        """Initialize Mostly AI SDK"""
        if not MOSTLY_AI_AVAILABLE:
            return False, "Mostly AI SDK not installed. Please install with: pip install mostlyai[local]"
        
        try:
            self.mostly = MostlyAI(local=True, local_port=8080)
            return True, "Mostly AI SDK initialized successfully!"
        except Exception as e:
            return False, f"Failed to initialize Mostly AI SDK: {str(e)}"
    
    
    def train_generator(self, data: pd.DataFrame, name: str, epochs: int = 10, max_training_time: int = 60, batch_size: int = 32, value_protection: bool = True) -> Tuple[bool, str]:
        """Train the synthetic data generator"""
        if not self.mostly:
            return False, "Mostly AI SDK not initialized"
        
        try:
            self.original_data = data
            train_config = {'tables':
                            [
                                {
                                    'name': name,
                                    'data': data,
                                    'tabular_model_configuration':
                                    {
                                        'max_epochs': epochs,
                                        'max_training_time': max_training_time,
                                        'value_protection': value_protection,
                                        'batch_size': batch_size
                                    }
                                }
                            ]
                           }
                            
            self.generator = self.mostly.train(
                config = train_config
            )
            return True, f"Generator trained successfully! Model: {name}"
        except Exception as e:
            return False, f"Training failed: {str(e)}"
    
    def generate_synthetic_data(self, size: int) -> Tuple[pd.DataFrame, str]:
        """Generate synthetic data"""
        if not self.generator:
            return None, "No trained generator available"
        
        try:
            synthetic_data = self.mostly.generate(self.generator, size=size)
            df = synthetic_data.data()
            return df, f"Generated {len(df)} synthetic records successfully!"
        except Exception as e:
            return None, f"Generation failed: {str(e)}"
    
    def get_quality_report(self) -> str:
        """Get quality assurance report"""
        if not self.generator:
            return "No trained generator available"
        
        try:
            report = self.generator.reports(display=False)
            return str(report)
        except Exception as e:
            return f"Failed to generate report: {str(e)}"
    
    def estimate_memory_usage(self, df: pd.DataFrame) -> str:
        """Estimate memory usage for the dataset"""
        if df is None or df.empty:
            return "No data to analyze"
        
        # Calculate approximate memory usage
        memory_mb = df.memory_usage(deep=True).sum() / (1024 * 1024)
        rows, cols = len(df), len(df.columns)
        
        # Estimate training memory (roughly 3-5x the data size)
        estimated_training_mb = memory_mb * 4
        
        status = "✅ Good" if memory_mb < 100 else "⚠️ Large" if memory_mb < 500 else "❌ Very Large"
        
        return f"""
**Memory Usage Estimate:**
- Data size: {memory_mb:.1f} MB
- Estimated training memory: {estimated_training_mb:.1f} MB
- Status: {status}
- Rows: {rows:,} | Columns: {cols}
        """.strip()

# Initialize the generator
generator = SyntheticDataGenerator()


def initialize_sdk() -> Tuple[str, str]:
    """Initialize the Mostly AI SDK"""
    success, message = generator.initialize_mostly_ai()
    status = "✅ Success" if success else "❌ Error"
    return status, message

def train_model(data: pd.DataFrame, model_name: str, epochs: int, max_training_time: int, batch_size: int, value_protection: bool) -> Tuple[str, str]:
    """Train the synthetic data generator"""
    if data is None or data.empty:
        return "❌ Error", "Please upload or create sample data first"
    
    success, message = generator.train_generator(data, model_name, epochs, max_training_time, batch_size, value_protection)
    status = "✅ Success" if success else "❌ Error"
    return status, message

def generate_data(size: int) -> Tuple[pd.DataFrame, str]:
    """Generate synthetic data"""
    if generator.generator is None:
        return None, "❌ Please train a model first"
    
    synthetic_df, message = generator.generate_synthetic_data(size)
    if synthetic_df is not None:
        status = "✅ Success"
    else:
        status = "❌ Error"
    
    return synthetic_df, f"{status} - {message}"

def get_quality_report() -> str:
    """Get quality report"""
    return generator.get_quality_report()

def create_comparison_plot(original_df: pd.DataFrame, synthetic_df: pd.DataFrame) -> go.Figure:
    """Create comparison plots between original and synthetic data"""
    if original_df is None or synthetic_df is None:
        return None
    
    # Select numeric columns for comparison
    numeric_cols = original_df.select_dtypes(include=[np.number]).columns.tolist()
    
    if not numeric_cols:
        return None
    
    # Create subplots
    n_cols = min(3, len(numeric_cols))
    n_rows = (len(numeric_cols) + n_cols - 1) // n_cols
    
    fig = make_subplots(
        rows=n_rows, 
        cols=n_cols,
        subplot_titles=numeric_cols[:n_rows*n_cols]
    )
    
    for i, col in enumerate(numeric_cols[:n_rows*n_cols]):
        row = i // n_cols + 1
        col_idx = i % n_cols + 1
        
        # Add original data histogram
        fig.add_trace(
            go.Histogram(
                x=original_df[col],
                name=f'Original {col}',
                opacity=0.7,
                nbinsx=20
            ),
            row=row, col=col_idx
        )
        
        # Add synthetic data histogram
        fig.add_trace(
            go.Histogram(
                x=synthetic_df[col],
                name=f'Synthetic {col}',
                opacity=0.7,
                nbinsx=20
            ),
            row=row, col=col_idx
        )
    
    fig.update_layout(
        title="Original vs Synthetic Data Comparison",
        height=300 * n_rows,
        showlegend=True
    )
    
    return fig

def download_csv(df: pd.DataFrame) -> str:
    """Convert DataFrame to CSV for download"""
    if df is None or df.empty:
        return None
    
    csv = df.to_csv(index=False)
    return csv

# Create the Gradio interface
def create_interface():
    with gr.Blocks(title="MOSTLY AI Synthetic Data Generator", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🎭 MOSTLY AI Synthetic Data Generator
        
        Generate high-quality synthetic data using the Mostly AI SDK. Upload your own CSV files to generate synthetic data that preserves the statistical properties of your original dataset.
        """)
        
        with gr.Tab("🚀 Quick Start"):
            gr.Markdown("### Initialize the SDK and upload your data")
            
            with gr.Row():
                with gr.Column():
                    init_btn = gr.Button("Initialize Mostly AI SDK", variant="primary")
                    init_status = gr.Textbox(label="Initialization Status", interactive=False)
                
                with gr.Column():
                    gr.Markdown("""
                    **Next Steps:**
                    1. Initialize the SDK (click button above)
                    2. Go to "Upload Data and Train Model" tab to upload your CSV file
                    3. Train a model on your data
                    4. Generate synthetic data
                    """)
        
        with gr.Tab("📊 Upload Data and Train Model"):
            gr.Markdown("### Upload your CSV file to generate synthetic data")
            
            gr.Markdown("""
            **📋 File Requirements:**
            - **Format:** CSV with header row
            - **Size:** Optimized for Hugging Face Spaces (2 vCPU, 16GB RAM)
            """)
            
            file_upload = gr.File(
                label="Upload CSV File",
                file_types=[".csv"],
                file_count="single"
            )
            
            uploaded_data = gr.Dataframe(label="Uploaded Data", interactive=False)
            
            memory_info = gr.Markdown(label="Memory Usage Info", visible=False)
            
            with gr.Row():
                with gr.Column():
                    model_name = gr.Textbox(
                        value="My Synthetic Model",
                        label="Model Name",
                        placeholder="Enter a name for your model"
                    )
                    epochs = gr.Slider(1, 200, value=100, step=1, label="Training Epochs")
                    max_training_time = gr.Slider(1, 1000, value=60, step=1, label="Maximum Training Time")
                    batch_size = gr.Slider(8, 1024, value=32, step=8, label="Training Batch Size")
                    value_protection = gr.Checkbox(label="Value Protection", info="Enable Value Protection")
                    train_btn = gr.Button("Train Model", variant="primary")
                
                with gr.Column():
                    train_status = gr.Textbox(label="Training Status", interactive=False)
                    quality_report = gr.Textbox(label="Quality Report", lines=10, interactive=False)
            
            get_report_btn = gr.Button("Get Quality Report", variant="secondary")
        
        with gr.Tab("🎲 Generate Data"):
            gr.Markdown("### Generate synthetic data from your trained model")
            
            with gr.Row():
                with gr.Column():
                    gen_size = gr.Slider(10, 1000, value=100, step=10, label="Number of Records to Generate")
                    generate_btn = gr.Button("Generate Synthetic Data", variant="primary")
                
                with gr.Column():
                    gen_status = gr.Textbox(label="Generation Status", interactive=False)
            
            synthetic_data = gr.Dataframe(label="Synthetic Data", interactive=False)
            
            with gr.Row():
                download_btn = gr.DownloadButton("Download CSV", variant="secondary")
                comparison_plot = gr.Plot(label="Data Comparison")
        
        # Event handlers
        init_btn.click(
            initialize_sdk,
            outputs=[init_status, init_status]
        )
        
        train_btn.click(
            train_model,
            inputs=[uploaded_data, model_name, epochs, max_training_time, batch_size, value_protection],
            outputs=[train_status, train_status]
        )
        
        get_report_btn.click(
            get_quality_report,
            outputs=[quality_report]
        )
        
        generate_btn.click(
            generate_data,
            inputs=[gen_size],
            outputs=[synthetic_data, gen_status]
        )
        
        # Update download button when synthetic data changes
        synthetic_data.change(
            download_csv,
            inputs=[synthetic_data],
            outputs=[download_btn]
        )
        
        # Create comparison plot when both datasets are available
        synthetic_data.change(
            create_comparison_plot,
            inputs=[uploaded_data, synthetic_data],
            outputs=[comparison_plot]
        )
        
        # Handle file upload with size and column limits
        def process_uploaded_file(file):
            if file is None:
                return None, "No file uploaded", gr.update(visible=False)
            
            try:
                # Read the CSV file
                df = pd.read_csv(file.name)
                
                # # Check column limit (max 20 columns)
                # if len(df.columns) > 20:
                #     return None, f"❌ Too many columns! Maximum allowed: 20, found: {len(df.columns)}. Please reduce the number of columns in your CSV file.", gr.update(visible=False)
                
                # # Check row limit (max 10,000 records)
                # if len(df) > 10000:
                #     return None, f"❌ Too many records! Maximum allowed: 10,000, found: {len(df)}. Please reduce the number of rows in your CSV file.", gr.update(visible=False)
                
                # # Check minimum requirements
                # if len(df) < 1000:
                #     return None, f"❌ Too few records! Minimum required: 1,000, found: {len(df)}. Please provide more data for training.", gr.update(visible=False)
                
                # if len(df.columns) < 2:
                #     return None, f"❌ Too few columns! Minimum required: 2, found: {len(df.columns)}. Please provide more columns for training.", gr.update(visible=False)
                
                # Success message with file info
                success_msg = f"✅ File uploaded successfully! {len(df)} rows × {len(df.columns)} columns"
                
                # Generate memory usage info
                memory_info = generator.estimate_memory_usage(df)
                
                return df, success_msg, gr.update(value=memory_info, visible=True)
                
            except Exception as e:
                return None, f"❌ Error reading file: {str(e)}", gr.update(visible=False)
        
        file_upload.change(
            process_uploaded_file,
            inputs=[file_upload],
            outputs=[uploaded_data, train_status, memory_info]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )