app.py

# Flask web application for CAMS air pollution visualization

import os
import json
import traceback
from pathlib import Path
import xarray as xr
import numpy as np

from datetime import datetime, timedelta
from werkzeug.utils import secure_filename
from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, send_file

# Import our custom modules
from data_processor import NetCDFProcessor, AuroraPredictionProcessor, analyze_netcdf_file
from plot_generator import IndiaMapPlotter
from interactive_plot_generator import InteractiveIndiaMapPlotter
from cams_downloader import CAMSDownloader
from constants import ALLOWED_EXTENSIONS, MAX_FILE_SIZE, COLOR_THEMES

# Aurora pipeline imports - with error handling for optional dependency
try:
    from aurora import Batch, Metadata, AuroraAirPollution, rollout
    from aurora_pipeline import AuroraPipeline
    AURORA_AVAILABLE = True
except ImportError as e:
    print(f"⚠️  Aurora model not available: {e}")
    AURORA_AVAILABLE = False

app = Flask(__name__)
app.secret_key = 'your-secret-key-change-this-in-production'  # Change this!
app.config['DEBUG'] = False  # Explicitly disable debug mode

# Add JSON filter for templates
import json
app.jinja_env.filters['tojson'] = json.dumps

# Configure upload settings
app.config['MAX_CONTENT_LENGTH'] = MAX_FILE_SIZE
app.config['UPLOAD_FOLDER'] = 'uploads'

# Initialize our services
downloader = CAMSDownloader()
plotter = IndiaMapPlotter()
interactive_plotter = InteractiveIndiaMapPlotter()

# Initialize Aurora pipeline if available
if AURORA_AVAILABLE:
    # Check if we're in development/local mode
    import socket
    hostname = socket.gethostname()
    is_local = any(local_indicator in hostname.lower() 
                   for local_indicator in ['local', 'macbook', 'laptop', 'desktop', 'dev'])
    
    # Force CPU mode for local development to avoid GPU requirements
    cpu_only = is_local or os.getenv('AURORA_CPU_ONLY', 'false').lower() == 'true'
    
    aurora_pipeline = AuroraPipeline(cpu_only=cpu_only)
    print(f"🔮 Aurora pipeline initialized ({'CPU-only' if cpu_only else 'GPU-enabled'} mode)")
else:
    aurora_pipeline = None
    print("⚠️  Aurora pipeline not available - missing dependencies")

# Ensure directories exist
for directory in ['uploads', 'downloads', 'plots', 'templates', 'static', 'predictions']:
    Path(directory).mkdir(exist_ok=True)


def allowed_file(filename):
    """Check if file extension is allowed"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def str_to_bool(value):
    """Convert string representation to boolean"""
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        return value.lower() in ('true', '1', 'yes', 'on')
    return bool(value)


@app.route('/')
def index():
    """Main page - file upload or date selection"""
    downloaded_files = downloader.list_downloaded_files()
    # List files in uploads and downloads/extracted
    upload_files = sorted(
        [f for f in Path(app.config['UPLOAD_FOLDER']).glob('*') if f.is_file()],
        key=lambda x: x.stat().st_mtime, reverse=True
    )
    extracted_files = sorted(
        [f for f in Path('downloads/extracted').glob('*') if f.is_file()],
        key=lambda x: x.stat().st_mtime, reverse=True
    )
    # Prepare for template: list of dicts with name and type
    recent_files = [
        {'name': f.name, 'type': 'upload'} for f in upload_files
    ] + [
        {'name': f.name, 'type': 'download'} for f in extracted_files
    ]
    current_date = datetime.now().strftime('%Y-%m-%d')
    return render_template(
        'index.html',
        downloaded_files=downloaded_files,
        cds_ready=downloader.is_client_ready(),
        current_date=current_date,
        recent_files=recent_files,
        aurora_available=AURORA_AVAILABLE
    )


@app.route('/upload', methods=['POST'])
def upload_file():
    """Handle file upload"""
    if 'file' not in request.files:
        flash('No file selected', 'error')
        return redirect(request.url)
    
    file = request.files['file']
    if file.filename == '':
        flash('No file selected', 'error')
        return redirect(request.url)
    
    if file and allowed_file(file.filename):
        try:
            filename = secure_filename(file.filename)
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f"{timestamp}_{filename}"
            filepath = Path(app.config['UPLOAD_FOLDER']) / filename
            
            file.save(str(filepath))
            flash(f'File uploaded successfully: {filename}', 'success')
            
            return redirect(url_for('analyze_file', filename=filename))
            
        except Exception as e:
            flash(f'Error uploading file: {str(e)}', 'error')
            return redirect(url_for('index'))
    else:
        flash('Invalid file type. Please upload .nc or .zip files.', 'error')
        return redirect(url_for('index'))


@app.route('/download_date', methods=['POST'])
def download_date():
    """Handle date-based download"""
    date_str = request.form.get('date')
    
    if not date_str:
        flash('Please select a date', 'error')
        return redirect(url_for('index'))
    
    # --- Backend Validation Logic ---
    try:
        selected_date = datetime.strptime(date_str, '%Y-%m-%d')
        start_date = datetime(2015, 1, 1)
        end_date = datetime.now()
        
        if not (start_date <= selected_date <= end_date):
            flash(f'Invalid date. Please select a date between {start_date.strftime("%Y-%m-%d")} and today.', 'error')
            return redirect(url_for('index'))
            
    except ValueError:
        flash('Invalid date format. Please use YYYY-MM-DD.', 'error')
        return redirect(url_for('index'))
    
    # --- End of Validation Logic ---
    
    if not downloader.is_client_ready():
        flash('CDS API not configured. Please check your environment variables or .cdsapirc file.', 'error')
        return redirect(url_for('index'))
    
    try:
        # Download CAMS data
        zip_path = downloader.download_cams_data(date_str)
        
        # Extract the files
        extracted_files = downloader.extract_cams_files(zip_path)
        
        flash(f'CAMS data downloaded successfully for {date_str}', 'success')
        
        # Analyze the extracted files
        if 'surface' in extracted_files:
            filename = Path(extracted_files['surface']).name
            return redirect(url_for('analyze_file', filename=filename, is_download='true'))
        elif 'atmospheric' in extracted_files:
            filename = Path(extracted_files['atmospheric']).name
            return redirect(url_for('analyze_file', filename=filename, is_download='true'))
        else:
            # Use the first available file
            first_file = list(extracted_files.values())[0]
            filename = Path(first_file).name
            return redirect(url_for('analyze_file', filename=filename, is_download='true'))
            
    except Exception as e:
        flash(f'Error downloading CAMS data: {str(e)}', 'error')
        return redirect(url_for('index'))


@app.route('/analyze/<filename>')
def analyze_file(filename):
    """Analyze uploaded file and show variable selection"""
    is_download_param = request.args.get('is_download', 'false')
    is_download = str_to_bool(is_download_param)
    
    try:
        # Determine file path
        if is_download:
            file_path = Path('downloads/extracted') / filename
        else:
            file_path = Path(app.config['UPLOAD_FOLDER']) / filename
        
        if not file_path.exists():
            flash('File not found', 'error')
            return redirect(url_for('index'))
        
        # Analyze the file
        analysis = analyze_netcdf_file(str(file_path))
        
        if not analysis['success']:
            flash(f'Error analyzing file: {analysis["error"]}', 'error')
            return redirect(url_for('index'))
        
        if analysis['total_variables'] == 0:
            flash('No air pollution variables found in the file', 'warning')
            return redirect(url_for('index'))
        
        # Process variables for template
        variables = []
        for var_name, var_info in analysis['detected_variables'].items():
            variables.append({
                'name': var_name,
                'display_name': var_info['name'],
                'type': var_info['type'],
                'units': var_info['units'],
                'shape': var_info['shape']
            })
        
        return render_template('variables.html',
                             filename=filename,
                             variables=variables,
                             color_themes=COLOR_THEMES,
                             is_download=is_download)
        
    except Exception as e:
        flash(f'Error analyzing file: {str(e)}', 'error')
        return redirect(url_for('index'))


@app.route('/get_pressure_levels/<filename>/<variable>')
def get_pressure_levels(filename, variable):
    """AJAX endpoint to get pressure levels for atmospheric variables"""
    try:
        is_download_param = request.args.get('is_download', 'false')
        is_download = str_to_bool(is_download_param)
        
        print(f"is_download: {is_download} (type: {type(is_download)})")

        # Determine file path
        if is_download:
            file_path = Path('downloads/extracted') / filename
            print("Using downloaded file path")
        else:
            file_path = Path(app.config['UPLOAD_FOLDER']) / filename
            print("Using upload file path")
        
        print(f"File path: {file_path}")

        processor = NetCDFProcessor(str(file_path))
        try:
            processor.load_dataset()
            processor.detect_variables()
            
            pressure_levels = processor.get_available_pressure_levels(variable)
        finally:
            processor.close()
        
        return jsonify({
            'success': True,
            'pressure_levels': pressure_levels
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


@app.route('/get_available_times/<filename>/<variable>')
def get_available_times(filename, variable):
    """AJAX endpoint to get available timestamps for a variable"""
    try:
        is_download_param = request.args.get('is_download', 'false')
        is_download = str_to_bool(is_download_param)
        
        # Determine file path
        if is_download:
            file_path = Path('downloads/extracted') / filename
        else:
            file_path = Path(app.config['UPLOAD_FOLDER']) / filename
        
        processor = NetCDFProcessor(str(file_path))
        try:
            processor.load_dataset()
            processor.detect_variables()
            
            available_times = processor.get_available_times(variable)
        finally:
            processor.close()
        
        # Format times for display
        formatted_times = []
        for i, time_val in enumerate(available_times):
            formatted_times.append({
                'index': i,
                'value': str(time_val),
                'display': time_val.strftime('%Y-%m-%d %H:%M') if hasattr(time_val, 'strftime') else str(time_val)
            })
        
        return jsonify({
            'success': True,
            'times': formatted_times
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })

@app.route('/visualize', methods=['POST'])
def visualize():
    """Generate and display the pollution map"""
    try:
        filename = request.form.get('filename')
        variable = request.form.get('variable')
        color_theme = request.form.get('color_theme', 'viridis')
        pressure_level = request.form.get('pressure_level')
        is_download_param = request.form.get('is_download', 'false')
        is_download = str_to_bool(is_download_param)
        
        if not filename or not variable:
            flash('Missing required parameters', 'error')
            return redirect(url_for('index'))
        
        # Determine file path
        if is_download:
            file_path = Path('downloads/extracted') / filename
        else:
            file_path = Path(app.config['UPLOAD_FOLDER']) / filename
        
        if not file_path.exists():
            flash('File not found', 'error')
            return redirect(url_for('index'))
        
        # Process the data
        processor = NetCDFProcessor(str(file_path))
        try:
            processor.load_dataset()
            processor.detect_variables()
            
            # Convert pressure level to float if provided
            pressure_level_val = None
            if pressure_level and pressure_level != 'None':
                try:
                    pressure_level_val = float(pressure_level)
                except ValueError:
                    pressure_level_val = None

            time_index_val = request.form.get('time_index')
            # Extract data
            data_values, metadata = processor.extract_data(
                variable, 
                time_index = int(time_index_val) if time_index_val and time_index_val != 'None' else 0,
                pressure_level=pressure_level_val
            )
            
            # Generate plot
            plot_path = plotter.create_india_map(
                data_values, 
                metadata, 
                color_theme=color_theme,
                save_plot=True
            )
            
        finally:
            # Always close the processor
            processor.close()
        
        if plot_path:
            plot_filename = Path(plot_path).name
            
            # Prepare metadata for display
            plot_info = {
                'variable': metadata.get('display_name', 'Unknown Variable'),
                'units': metadata.get('units', ''),
                'shape': str(metadata.get('shape', 'Unknown')),
                'pressure_level': metadata.get('pressure_level'),
                'color_theme': COLOR_THEMES.get(color_theme, color_theme),
                'generated_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'data_range': {
                    'min': float(f"{data_values.min():.3f}") if hasattr(data_values, 'min') and not data_values.min() is None else 0,
                    'max': float(f"{data_values.max():.3f}") if hasattr(data_values, 'max') and not data_values.max() is None else 0,
                    'mean': float(f"{data_values.mean():.3f}") if hasattr(data_values, 'mean') and not data_values.mean() is None else 0
                }
            }
            
            print(f"Plot info prepared: {plot_info}")
            
            return render_template('plot.html',
                                    plot_filename=plot_filename,
                                    plot_info=plot_info)
        else:
            flash('Error generating plot', 'error')
            return redirect(url_for('index'))
            
    except Exception as e:
        flash(f'Error creating visualization: {str(e)}', 'error')
        print(f"Full error: {traceback.format_exc()}")
        return redirect(url_for('index'))
    
@app.route('/visualize_interactive', methods=['POST'])
def visualize_interactive():
    """Generate and display the interactive pollution map"""
    try:
        filename = request.form.get('filename')
        variable = request.form.get('variable')
        color_theme = request.form.get('color_theme', 'viridis')
        pressure_level = request.form.get('pressure_level')
        is_download_param = request.form.get('is_download', 'false')
        is_download = str_to_bool(is_download_param)
        
        if not filename or not variable:
            flash('Missing required parameters', 'error')
            return redirect(url_for('index'))
        
        # Determine file path
        if is_download:
            file_path = Path('downloads/extracted') / filename
        else:
            file_path = Path(app.config['UPLOAD_FOLDER']) / filename
        
        if not file_path.exists():
            flash('File not found', 'error')
            return redirect(url_for('index'))
        
        # Process the data
        processor = NetCDFProcessor(str(file_path))
        try:
            processor.load_dataset()
            processor.detect_variables()
            
            # Convert pressure level to float if provided
            pressure_level_val = None
            if pressure_level and pressure_level != 'None':
                try:
                    pressure_level_val = float(pressure_level)
                except ValueError:
                    pressure_level_val = None

            time_index_val = request.form.get('time_index')
            # Extract data
            data_values, metadata = processor.extract_data(
                variable, 
                time_index = int(time_index_val) if time_index_val and time_index_val != 'None' else 0,
                pressure_level=pressure_level_val
            )
            
            # Generate interactive plot
            result = interactive_plotter.create_india_map(
                data_values, 
                metadata, 
                color_theme=color_theme,
                save_plot=True
            )
        
        finally:
            # Always close the processor
            processor.close()
        
        if result and result.get('html_content'):
            # Prepare metadata for display
            plot_info = {
                'variable': metadata.get('display_name', 'Unknown Variable'),
                'units': metadata.get('units', ''),
                'shape': str(metadata.get('shape', 'Unknown')),
                'pressure_level': metadata.get('pressure_level'),
                'color_theme': COLOR_THEMES.get(color_theme, color_theme),
                'generated_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'data_range': {
                    'min': float(f"{data_values.min():.3f}") if hasattr(data_values, 'min') and not data_values.min() is None else 0,
                    'max': float(f"{data_values.max():.3f}") if hasattr(data_values, 'max') and not data_values.max() is None else 0,
                    'mean': float(f"{data_values.mean():.3f}") if hasattr(data_values, 'mean') and not data_values.mean() is None else 0
                },
                'is_interactive': True,
                'html_path': result.get('html_path'),
                'png_path': result.get('png_path')
            }
            
            return render_template('interactive_plot.html',
                                    plot_html=result['html_content'],
                                    plot_info=plot_info)
        else:
            flash('Error generating interactive plot', 'error')
            return redirect(url_for('index'))
            
    except Exception as e:
        flash(f'Error creating interactive visualization: {str(e)}', 'error')
        print(f"Full error: {traceback.format_exc()}")
        return redirect(url_for('index'))
    
@app.route('/plot/<filename>')
def serve_plot(filename):
    """Serve plot images"""
    try:
        plot_path = Path('plots') / filename
        if plot_path.exists():
            # Determine mimetype based on file extension
            if filename.lower().endswith('.html'):
                return send_file(str(plot_path), mimetype='text/html', as_attachment=True)
            elif filename.lower().endswith('.jpg') or filename.lower().endswith('.jpeg'):
                mimetype = 'image/jpeg'
                return send_file(str(plot_path), mimetype=mimetype)
            elif filename.lower().endswith('.png'):
                mimetype = 'image/png'
                return send_file(str(plot_path), mimetype=mimetype)
            else:
                return send_file(str(plot_path), as_attachment=True)
        else:
            flash('Plot not found', 'error')
            return redirect(url_for('index'))
    except Exception as e:
        flash(f'Error serving plot: {str(e)}', 'error')
        return redirect(url_for('index'))


@app.route('/gallery')
def gallery():
    """Display gallery of all saved plots"""
    try:
        plots_dir = Path('plots')
        plots_dir.mkdir(exist_ok=True)
        
        # Get all plot files
        plot_files = []
        
        # Static plots (PNG/JPG)
        for ext in ['*.png', '*.jpg', '*.jpeg']:
            for plot_file in plots_dir.glob(ext):
                if plot_file.is_file():
                    # Parse filename to extract metadata
                    filename = plot_file.name
                    file_info = {
                        'filename': filename,
                        'path': str(plot_file),
                        'type': 'static',
                        'size': plot_file.stat().st_size,
                        'created': datetime.fromtimestamp(plot_file.stat().st_mtime),
                        'extension': plot_file.suffix.lower()
                    }
                    
                    # Try to parse metadata from filename
                    try:
                        # Example: PM2_5_India_viridis_20200824_1200.png or PM2_5_India_850hPa_viridis_20200824_1200.png
                        name_parts = filename.replace(plot_file.suffix, '').split('_')
                        if len(name_parts) >= 4:
                            # Extract variable name (everything before _India)
                            var_parts = []
                            for i, part in enumerate(name_parts):
                                if part == 'India':
                                    break
                                var_parts.append(part)
                            
                            file_info['variable'] = '_'.join(var_parts).replace('_', '.')
                            file_info['region'] = 'India'
                            file_info['plot_type'] = 'Static'
                            
                            # Check if there's pressure level (contains 'hPa')
                            pressure_level = None
                            theme_color = 'Unknown'
                            for part in name_parts:
                                if 'hPa' in part:
                                    pressure_level = part
                                elif part not in var_parts and part != 'India' and not part.isdigit():
                                    # This is likely the color theme
                                    theme_color = part
                                    break
                            
                            file_info['pressure_level'] = pressure_level
                            file_info['theme'] = theme_color
                    except:
                        file_info['variable'] = 'Unknown'
                        file_info['region'] = 'Unknown'
                        file_info['theme'] = 'Unknown'
                        file_info['pressure_level'] = None
                        file_info['plot_type'] = 'Static'
                    
                    plot_files.append(file_info)
        
        # Interactive plots (HTML)
        for plot_file in plots_dir.glob('*.html'):
            if plot_file.is_file():
                filename = plot_file.name
                file_info = {
                    'filename': filename,
                    'path': str(plot_file),
                    'type': 'interactive',
                    'size': plot_file.stat().st_size,
                    'created': datetime.fromtimestamp(plot_file.stat().st_mtime),
                    'extension': '.html'
                }
                
                # Try to parse metadata from filename
                try:
                    # Example: PM2_5_India_interactive_viridis_20200824_1200.html or PM2_5_India_interactive_850hPa_viridis_20200824_1200.html
                    name_parts = filename.replace('.html', '').split('_')
                    if len(name_parts) >= 5:
                        # Extract variable name (everything before _India)
                        var_parts = []
                        for i, part in enumerate(name_parts):
                            if part == 'India':
                                break
                            var_parts.append(part)
                        
                        file_info['variable'] = '_'.join(var_parts).replace('_', '.')
                        file_info['region'] = 'India'
                        file_info['plot_type'] = 'Interactive'
                        
                        # Check if there's pressure level (contains 'hPa')
                        pressure_level = None
                        theme_color = 'Unknown'
                        for part in name_parts:
                            if 'hPa' in part:
                                pressure_level = part
                            elif part not in var_parts and part not in ['India', 'interactive'] and not part.isdigit():
                                # This is likely the color theme
                                theme_color = part
                                break
                        
                        file_info['pressure_level'] = pressure_level
                        file_info['theme'] = theme_color
                except:
                    file_info['variable'] = 'Unknown'
                    file_info['region'] = 'Unknown'
                    file_info['theme'] = 'Unknown'
                    file_info['pressure_level'] = None
                    file_info['plot_type'] = 'Interactive'
                
                plot_files.append(file_info)
        
        # Sort by creation time (newest first)
        plot_files.sort(key=lambda x: x['created'], reverse=True)
        
        # Group by type for display
        static_plots = [p for p in plot_files if p['type'] == 'static']
        interactive_plots = [p for p in plot_files if p['type'] == 'interactive']
        
        return render_template('gallery.html',
                             static_plots=static_plots,
                             interactive_plots=interactive_plots,
                             total_plots=len(plot_files))
        
    except Exception as e:
        flash(f'Error loading gallery: {str(e)}', 'error')
        return redirect(url_for('index'))


@app.route('/view_interactive/<filename>')
def view_interactive_plot(filename):
    """View an interactive plot from the gallery"""
    try:
        plot_path = Path('plots') / filename
        if not plot_path.exists() or not filename.endswith('.html'):
            flash('Interactive plot not found', 'error')
            return redirect(url_for('gallery'))
        
        # Read the HTML content
        with open(plot_path, 'r', encoding='utf-8') as f:
            html_content = f.read()
        
        # Create plot info from filename parsing
        plot_info = {
            'variable': 'Unknown',
            'pressure_level': None,
            'theme': 'Unknown',
            'plot_type': 'Interactive',
            'generated_time': datetime.fromtimestamp(plot_path.stat().st_mtime).strftime('%Y-%m-%d %H:%M:%S'),
            'is_interactive': True,
            'filename': filename
        }
        
        # Try to parse metadata from filename
        try:
            # Example: PM2_5_India_interactive_850hPa_viridis_20200824_1200.html
            name_parts = filename.replace('.html', '').split('_')
            if len(name_parts) >= 5:
                # Extract variable name (everything before _India)
                var_parts = []
                for i, part in enumerate(name_parts):
                    if part == 'India':
                        break
                    var_parts.append(part)
                
                plot_info['variable'] = '_'.join(var_parts).replace('_', '.')
                
                # Check if there's pressure level (contains 'hPa')
                pressure_level = None
                theme_color = 'Unknown'
                for part in name_parts:
                    if 'hPa' in part:
                        pressure_level = part
                    elif part not in var_parts and part not in ['India', 'interactive'] and not part.isdigit():
                        # This is likely the color theme
                        theme_color = part
                        break
                
                plot_info['pressure_level'] = pressure_level
                plot_info['theme'] = theme_color
        except:
            pass  # Keep defaults
        
        return render_template('view_interactive.html',
                             plot_html=html_content,
                             plot_info=plot_info)
        
    except Exception as e:
        flash(f'Error viewing plot: {str(e)}', 'error')
        return redirect(url_for('gallery'))


@app.route('/delete_plot/<filename>', methods=['DELETE'])
def delete_plot(filename):
    """Delete a specific plot file"""
    print(f"🗑️  DELETE request received for: {filename}")  # Debug logging
    try:
        # Security check: ensure filename is safe
        if not filename or '..' in filename or '/' in filename:
            print(f"❌ Invalid filename: {filename}")
            return jsonify({'success': False, 'error': 'Invalid filename'}), 400
        
        plot_path = Path('plots') / filename
        print(f"🔍 Looking for file at: {plot_path}")
        
        # Check if file exists
        if not plot_path.exists():
            print(f"❌ File not found: {plot_path}")
            return jsonify({'success': False, 'error': 'File not found'}), 404
        
        # Check if it's a valid plot file
        allowed_extensions = ['.png', '.jpg', '.jpeg', '.html']
        if plot_path.suffix.lower() not in allowed_extensions:
            print(f"❌ Invalid file type: {filename}")
            return jsonify({'success': False, 'error': 'Invalid file type'}), 400
        
        # Delete the file
        plot_path.unlink()
        print(f"✅ Successfully deleted: {filename}")
        
        return jsonify({
            'success': True, 
            'message': f'Plot {filename} deleted successfully'
        })
        
    except Exception as e:
        print(f"💥 Error deleting file {filename}: {str(e)}")
        return jsonify({
            'success': False, 
            'error': f'Failed to delete plot: {str(e)}'
        }), 500


@app.route('/cleanup')
def cleanup():
    """Clean up old files"""
    try:
        # Clean up old plots (older than 24 hours)
        cutoff_time = datetime.now() - timedelta(hours=24)
        cleaned_count = 0
        
        for plot_file in Path('plots').glob('*.png'):
            if plot_file.stat().st_mtime < cutoff_time.timestamp():
                plot_file.unlink()
                cleaned_count += 1
        
        for plot_file in Path('plots').glob('*.jpg'):
            if plot_file.stat().st_mtime < cutoff_time.timestamp():
                plot_file.unlink()
                cleaned_count += 1
        
        flash(f'Cleaned up {cleaned_count} old plot files', 'success')
        return redirect(url_for('index'))
        
    except Exception as e:
        print(f"Error during cleanup: {str(e)}")
        flash('Error during cleanup', 'error')
        return redirect(url_for('index'))


@app.route('/health')
def health_check():
    """Health check endpoint for monitoring"""
    return jsonify({
        'status': 'healthy',
        'timestamp': datetime.now().isoformat(),
        'cds_ready': downloader.is_client_ready(),
        'aurora_available': AURORA_AVAILABLE
    })


@app.route('/api/aurora_status')
def aurora_status():
    """API endpoint to check Aurora readiness and get system info"""
    status = {
        'available': AURORA_AVAILABLE,
        'cpu_only': False,
        'estimated_time': {
            'cpu': {'1_step': 5, '2_steps': 10},
            'gpu': {'4_steps': 3, '6_steps': 4, '10_steps': 6}
        }
    }
    
    if AURORA_AVAILABLE and aurora_pipeline:
        status['cpu_only'] = getattr(aurora_pipeline, 'cpu_only', False)
    
    return jsonify(status)


# Aurora ML Prediction Routes
@app.route('/aurora_predict', methods=['GET', 'POST'])
def aurora_predict():
    """Aurora prediction form and handler with enhanced step selection"""
    if not AURORA_AVAILABLE:
        flash('Aurora model is not available. Please install required dependencies.', 'error')
        return redirect(url_for('index'))
    
    if request.method == 'GET':
        current_date = datetime.now().strftime('%Y-%m-%d')
        # Get list of existing prediction runs
        existing_runs = AuroraPipeline.list_prediction_runs() if hasattr(AuroraPipeline, 'list_prediction_runs') else []
        return render_template('aurora_predict.html', 
                               current_date=current_date,
                               existing_runs=existing_runs)
    
    # POST: Run the pipeline
    date_str = request.form.get('date')
    steps = int(request.form.get('steps', 2))  # Default to 2 steps
    
    # Validate steps (1-4 allowed, each representing 12 hours)
    if steps < 1 or steps > 4:
        flash('Number of steps must be between 1 and 4 (each step = 12 hours)', 'error')
        return redirect(url_for('aurora_predict'))
    
    # Limit steps for local/CPU execution
    if hasattr(aurora_pipeline, 'cpu_only') and aurora_pipeline.cpu_only:
        max_cpu_steps = 2
        if steps > max_cpu_steps:
            steps = max_cpu_steps
            flash(f'Steps reduced to {steps} for CPU mode optimization', 'info')
    
    if not date_str:
        flash('Please select a valid date.', 'error')
        return redirect(url_for('aurora_predict'))

    try:
        print(f"🚀 Starting Aurora prediction pipeline for {date_str}")
        print(f"📊 Requested {steps} forward steps ({steps * 12} hours coverage)")
        
        # 1. Download CAMS data for the selected date (if not already available)
        print("📥 Step 1/5: Checking/downloading CAMS atmospheric data...")
        try:
            zip_path = downloader.download_cams_data(date_str)
        except Exception as e:
            error_msg = f"Failed to download CAMS data: {str(e)}"
            if "error response" in str(e).lower():
                error_msg += " The CAMS API may have returned an error. Please try a different date or check your CDS API credentials."
            elif "zip" in str(e).lower():
                error_msg += " The downloaded file is corrupted. Please try again."
            flash(error_msg, 'error')
            print(f"❌ Download error: {traceback.format_exc()}")
            return redirect(url_for('aurora_predict'))
        
        try:
            extracted_files = downloader.extract_cams_files(zip_path)
            print("✅ CAMS data downloaded and extracted")
        except Exception as e:
            error_msg = f"Failed to extract CAMS data: {str(e)}"
            if "not a zip file" in str(e).lower():
                error_msg += " The downloaded file appears to be corrupted or is an error response from the CAMS API."
            elif "html" in str(e).lower() or "error" in str(e).lower():
                error_msg += " The CAMS API returned an error page instead of data."
            flash(error_msg, 'error')
            print(f"❌ Extraction error: {traceback.format_exc()}")
            return redirect(url_for('aurora_predict'))
        
        # 2. Run enhanced Aurora pipeline
        print("🔮 Step 2/5: Running enhanced Aurora ML pipeline...")
        
        try:
            # Use the enhanced pipeline method
            run_metadata = aurora_pipeline.run_aurora_prediction_pipeline(
                date_str=date_str,
                Batch=Batch,
                Metadata=Metadata,
                AuroraAirPollution=AuroraAirPollution,
                rollout=rollout,
                steps=steps
            )
            
            print("✅ Aurora predictions completed successfully")
            
            # Redirect to aurora variables page
            run_dir_name = run_metadata['run_directory'].split('/')[-1]
            
            flash(f'🔮 Aurora predictions generated successfully for {date_str} ({steps} steps, {steps * 12}h coverage)', 'success')
            return redirect(url_for('aurora_variables', run_dir=run_dir_name))
            
        except Exception as e:
            error_msg = f"Aurora model execution failed: {str(e)}"
            if "map_location" in str(e):
                error_msg += " This appears to be a compatibility issue with the Aurora model version."
            elif "checkpoint" in str(e).lower():
                error_msg += " Failed to load the Aurora model. Please check if the model files are properly installed."
            elif "memory" in str(e).lower() or "cuda" in str(e).lower():
                error_msg += " Insufficient memory or GPU issues. Try reducing the number of prediction steps."
            flash(error_msg, 'error')
            print(f"❌ Aurora model error: {traceback.format_exc()}")
            return redirect(url_for('aurora_predict'))
    
    except Exception as e:
        # Catch-all for any other unexpected errors
        error_msg = f'Unexpected error in Aurora pipeline: {str(e)}'
        flash(error_msg, 'error')
        print(f"❌ Unexpected Aurora pipeline error: {traceback.format_exc()}")
        return redirect(url_for('aurora_predict'))
    
    except Exception as e:
        # Catch-all for any other unexpected errors
        error_msg = f'Unexpected error in Aurora pipeline: {str(e)}'
        flash(error_msg, 'error')
        print(f"❌ Unexpected Aurora pipeline error: {traceback.format_exc()}")
        return redirect(url_for('aurora_predict'))


@app.route('/visualize_prediction/<path:filename>', methods=['GET', 'POST'])
def visualize_prediction(filename):
    """Aurora prediction visualization with step, variable, and pressure level selection"""
    # Handle both old and new filename formats
    if filename.endswith('.nc'):
        file_path = Path('predictions') / filename
    else:
        # Try to find the prediction file in the run directory
        run_dir = Path('predictions') / filename
        if run_dir.is_dir():
            # Look for the .nc file in the directory
            nc_files = list(run_dir.glob("*.nc"))
            if nc_files:
                file_path = nc_files[0]
            else:
                flash('No prediction file found in run directory', 'error')
                return redirect(url_for('index'))
        else:
            file_path = Path('predictions') / filename
    
    if not file_path.exists():
        flash('Prediction file not found', 'error')
        return redirect(url_for('index'))

    try:
        ds = xr.open_dataset(file_path)
        
        # Get all variables and separate surface from atmospheric
        all_variables = list(ds.data_vars.keys())
        surface_vars = []
        atmospheric_vars = []
        
        for var in all_variables:
            if 'pressure_level' in ds[var].dims:
                atmospheric_vars.append(var)
            else:
                surface_vars.append(var)
        
        # Get steps and pressure levels
        steps = list(range(len(ds['step']))) if 'step' in ds else [0]
        pressure_levels = list(ds['pressure_level'].values) if 'pressure_level' in ds else []
        
        # Handle form submission
        if request.method == 'POST':
            selected_step = int(request.form.get('step', 0))
            var_name = request.form.get('variable')
            pressure_level = request.form.get('pressure_level')
            color_theme = request.form.get('color_theme', 'viridis')
            plot_type = request.form.get('plot_type', 'static')
        else:
            selected_step = 0
            var_name = surface_vars[0] if surface_vars else all_variables[0] if all_variables else None
            pressure_level = None
            color_theme = 'viridis'
            plot_type = 'static'
            
        if not var_name or var_name not in all_variables:
            flash('Invalid variable selected', 'error')
            return redirect(url_for('index'))

        # Validate step
        if selected_step < 0 or selected_step >= len(steps):
            selected_step = 0

        return render_template(
            'aurora_variables.html',
            filename=filename,
            file_path=str(file_path),
            surface_vars=surface_vars,
            atmospheric_vars=atmospheric_vars,
            steps=steps,
            pressure_levels=pressure_levels,
            selected_step=selected_step,
            selected_variable=var_name,
            selected_pressure_level=pressure_level,
            color_theme=color_theme,
            plot_type=plot_type,
            color_themes=COLOR_THEMES
        )
            
    except Exception as e:
        flash(f'Error processing prediction file: {str(e)}', 'error')
        print(f"❌ Prediction visualization error: {traceback.format_exc()}")
        return redirect(url_for('index'))


@app.route('/generate_aurora_plot', methods=['POST'])
def generate_aurora_plot():
    """Generate plot from Aurora prediction data"""
    try:
        file_path = request.form.get('file_path')
        step = int(request.form.get('step', 0))
        var_name = request.form.get('variable')
        pressure_level = request.form.get('pressure_level')
        color_theme = request.form.get('color_theme', 'viridis')
        plot_type = request.form.get('plot_type', 'static')
        
        if not file_path or not var_name:
            flash('Missing required parameters', 'error')
            return redirect(url_for('index'))
        
        # Open dataset
        ds = xr.open_dataset(file_path)
        
        # Get data for the selected variable and step
        data = ds[var_name]
        
        # Handle different dimensions
        if 'step' in data.dims:
            data = data.isel(step=step)
        
        if pressure_level and 'pressure_level' in data.dims:
            pressure_level = float(pressure_level)
            data = data.sel(pressure_level=pressure_level, method='nearest')
        
        # Convert to numpy
        data_to_plot = data.values
        
        # Get coordinates
        lats = ds['lat'].values if 'lat' in ds else ds['latitude'].values
        lons = ds['lon'].values if 'lon' in ds else ds['longitude'].values
        
        # Prepare metadata
        from constants import NETCDF_VARIABLES
        var_info = NETCDF_VARIABLES.get(var_name, {})
        display_name = var_info.get('name', var_name)
        units = ds[var_name].attrs.get('units', var_info.get('units', ''))
        
        hours_from_start = (step + 1) * 12
        step_time_str = f"T+{hours_from_start}h (Step {step + 1})"
        
        metadata = {
            'variable_name': var_name,
            'display_name': display_name,
            'units': units,
            'lats': lats,
            'lons': lons,
            'pressure_level': pressure_level if pressure_level else None,
            'timestamp_str': step_time_str,
        }
        
        # Generate plot based on type
        if plot_type == 'interactive':
            # Generate interactive plot
            plot_result = interactive_plotter.create_india_map(
                data_to_plot,
                metadata,
                color_theme=color_theme,
                save_plot=True,
                custom_title=f"Aurora Prediction: {display_name} ({step_time_str})"
            )
            
            if plot_result and plot_result.get('html_path'):
                plot_filename = Path(plot_result['html_path']).name
                return render_template(
                    'interactive_plot.html',
                    plot_filename=plot_filename,
                    var_name=var_name,
                    pressure_level=pressure_level,
                    metadata=metadata
                )
        else:
            # Generate static plot
            plot_path = plotter.create_india_map(
                data_to_plot,
                metadata,
                color_theme=color_theme,
                save_plot=True,
                custom_title=f"Aurora Prediction: {display_name} ({step_time_str})"
            )
            
            if plot_path:
                plot_filename = Path(plot_path).name
                return render_template(
                    'plot.html',
                    plot_filename=plot_filename,
                    var_name=var_name,
                    pressure_level=pressure_level,
                    metadata=metadata
                )
        
        flash('Error generating plot', 'error')
        return redirect(url_for('index'))
        
    except Exception as e:
        flash(f'Error generating plot: {str(e)}', 'error')
        print(f"❌ Plot generation error: {traceback.format_exc()}")
        return redirect(url_for('index'))


@app.route('/aurora_plot', methods=['POST'])
def aurora_plot():
    """Generate plot from Aurora prediction variables"""
    if not AURORA_AVAILABLE:
        flash('Aurora model is not available.', 'error')
        return redirect(url_for('index'))
    
    try:
        run_dir = request.form.get('run_dir')
        step = request.form.get('step')
        variable = request.form.get('variable')
        pressure_level = request.form.get('pressure_level')
        color_theme = request.form.get('color_theme', 'viridis')
        plot_type = request.form.get('plot_type', 'static')
        
        if not all([run_dir, step, variable]):
            flash('Missing required parameters', 'error')
            return redirect(url_for('aurora_variables', run_dir=run_dir))
        
        # Find the filename for this step
        run_path = Path('predictions') / run_dir
        step_files = list(run_path.glob(f'*_step{int(step):02d}_*.nc'))
        
        if not step_files:
            flash(f'No file found for step {step}', 'error')
            return redirect(url_for('aurora_variables', run_dir=run_dir))
        
        file_path = step_files[0]  # Take the first match
        filename = file_path.name
        
        if not file_path.exists():
            flash('Prediction file not found', 'error')
            return redirect(url_for('aurora_variables', run_dir=run_dir))
        
        # Use Aurora prediction processor
        processor = AuroraPredictionProcessor(str(file_path))
        
        try:
            file_info = analyze_netcdf_file(str(file_path))
            
            var_info = file_info['detected_variables'].get(variable)
            if not var_info:
                flash('Variable not found in file', 'error')
                return redirect(url_for('aurora_variables', run_dir=run_dir))
            
            # Extract data using Aurora processor with step=0 (single timestep files)
            if var_info.get('type') == 'atmospheric' and pressure_level:
                pressure_level = float(pressure_level)
                data, metadata = processor.extract_variable_data(variable, pressure_level=pressure_level, step=0)
            else:
                data, metadata = processor.extract_variable_data(variable, step=0)
            

            
            # Prepare plot_info for templates
            plot_info = {
                'variable': metadata.get('display_name', 'Unknown Variable'),
                'units': metadata.get('units', ''),
                'shape': str(data.shape),
                'pressure_level': metadata.get('pressure_level'),
                'color_theme': COLOR_THEMES.get(color_theme, color_theme),
                'generated_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'data_range': {
                    'min': float(f"{data.min():.3f}") if hasattr(data, 'min') and data.min() is not None else 0,
                    'max': float(f"{data.max():.3f}") if hasattr(data, 'max') and data.max() is not None else 0,
                    'mean': float(f"{data.mean():.3f}") if hasattr(data, 'mean') and data.mean() is not None else 0
                },
                'timestamp': metadata.get('timestamp_str', 'Unknown Time'),
                'source': metadata.get('source', 'Aurora Prediction')
            }
            
            if plot_type == 'interactive':
                plot_result = interactive_plotter.create_india_map(
                    data, metadata, color_theme=color_theme, save_plot=True
                )
                if plot_result and plot_result.get('html_path'):
                    plot_filename = Path(plot_result['html_path']).name
                    return render_template('view_interactive.html', 
                                         plot_filename=plot_filename,
                                         metadata=metadata,
                                         plot_info=plot_info)
            else:
                plot_path = plotter.create_india_map(
                    data, metadata, color_theme=color_theme, save_plot=True
                )
                if plot_path:
                    plot_filename = Path(plot_path).name
                    return render_template('plot.html', 
                                         plot_filename=plot_filename,
                                         metadata=metadata,
                                         plot_info=plot_info,
                                         filename=filename)
            
            flash('Error generating plot', 'error')
            return redirect(url_for('aurora_variables', run_dir=run_dir))
            
        finally:
            processor.close()
        
    except Exception as e:
        flash(f'Error generating Aurora plot: {str(e)}', 'error')
        return redirect(url_for('index'))

@app.route('/download_prediction_netcdf/<path:filename>')
def download_prediction_netcdf(filename):
    """Download the Aurora prediction NetCDF file"""
    # Handle both old and new filename formats
    if filename.endswith('.nc'):
        file_path = Path('predictions') / filename
    else:
        # Try to find the prediction file in the run directory
        run_dir = Path('predictions') / filename
        if run_dir.is_dir():
            nc_files = list(run_dir.glob("*.nc"))
            if nc_files:
                file_path = nc_files[0]
                filename = file_path.name
            else:
                flash('Prediction file not found', 'error')
                return redirect(url_for('index'))
        else:
            file_path = Path('predictions') / filename
    
    if not file_path.exists():
        flash('Prediction file not found', 'error')
        return redirect(url_for('index'))
    return send_file(str(file_path), as_attachment=True, download_name=filename)


@app.errorhandler(413)
def too_large(e):
    """Handle file too large error"""
    flash('File too large. Maximum size is 500MB.', 'error')
    return redirect(url_for('index'))


@app.route('/api/aurora_step_variables/<run_dir>/<int:step>')
def get_aurora_step_variables(run_dir, step):
    """Get variables and pressure levels for a specific Aurora prediction step"""
    if not AURORA_AVAILABLE:
        return jsonify({'error': 'Aurora model not available'}), 400
    
    try:
        # Find the file for this step
        run_path = Path('predictions') / run_dir
        step_files = list(run_path.glob(f'*_step{step:02d}_*.nc'))
        
        if not step_files:
            return jsonify({'error': f'No file found for step {step}'}), 404
        
        file_path = step_files[0]
        
        # Load and analyze the file using the same method as regular CAMS files
        file_info = analyze_netcdf_file(str(file_path))
        
        if not file_info['success']:
            return jsonify({'error': f'Failed to analyze file: {file_info.get("error", "Unknown error")}'}), 500
        
        surface_vars = []
        atmos_vars = []
        pressure_levels = []
        
        # Extract variables from detected_variables
        for var_name, var_info in file_info['detected_variables'].items():
            if var_info['type'] == 'surface':
                surface_vars.append(var_name)
            elif var_info['type'] == 'atmospheric':
                atmos_vars.append(var_name)
        
        # Get pressure levels from the first atmospheric variable
        ds = xr.open_dataset(file_path)
        if 'pressure_level' in ds.coords:
            pressure_levels = ds.pressure_level.values.tolist()
        ds.close()
        
        return jsonify({
            'surface_vars': surface_vars,
            'atmos_vars': atmos_vars,
            'pressure_levels': pressure_levels,
            'filename': file_path.name
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/aurora_variables/<run_dir>')
def aurora_variables(run_dir):
    """Show Aurora prediction variables selection page similar to variables.html"""
    if not AURORA_AVAILABLE:
        flash('Aurora model is not available.', 'error')
        return redirect(url_for('index'))
    
    try:
        # Get prediction files from run directory
        run_path = Path('predictions') / run_dir
        if not run_path.exists():
            flash(f'Prediction run not found: {run_path}', 'error')
            return redirect(url_for('index'))
        
        # Find all prediction files in the directory
        pred_files = sorted(run_path.glob('*.nc'))
        if not pred_files:
            flash('No prediction files found in run', 'error')
            return redirect(url_for('index'))
        
        # Get step numbers and filenames
        steps_data = []
        for file_path in pred_files:
            filename = file_path.name
            # Extract step number from filename
            if 'step' in filename:
                try:
                    step_part = filename.split('step')[1].split('_')[0]
                    step_num = int(step_part)
                    steps_data.append({
                        'step': step_num,
                        'filename': filename,
                        'forecast_hours': step_num * 12
                    })
                except:
                    pass
        
        steps_data.sort(key=lambda x: x['step'])
        
        # Get variables from the first file
        first_file = pred_files[0]
        ds = xr.open_dataset(first_file)
        
        # Separate surface and atmospheric variables
        surface_vars = []
        atmos_vars = []
        pressure_levels = []
        
        for var_name in ds.data_vars:
            if len(ds[var_name].dims) == 2:  # lat, lon
                surface_vars.append(var_name)
            elif len(ds[var_name].dims) == 3:  # pressure_level, lat, lon
                atmos_vars.append(var_name)
        
        if 'pressure_level' in ds.coords:
            pressure_levels = ds.pressure_level.values.tolist()
        
        ds.close()
        
        return render_template('aurora_variables.html',
                             run_dir=run_dir,
                             steps_data=steps_data,
                             surface_vars=surface_vars,
                             atmos_vars=atmos_vars,
                             pressure_levels=pressure_levels,
                             color_themes=COLOR_THEMES)
                             
    except Exception as e:
        flash(f'Error loading Aurora variables: {str(e)}', 'error')
        return redirect(url_for('index'))

@app.route('/prediction_runs')
def prediction_runs():
    """Browse available Aurora prediction runs"""
    if not AURORA_AVAILABLE:
        flash('Aurora model is not available.', 'error')
        return redirect(url_for('index'))
    
    try:
        runs = AuroraPipeline.list_prediction_runs()
        return render_template('prediction_runs.html', runs=runs)
    except Exception as e:
        flash(f'Error listing prediction runs: {str(e)}', 'error')
        return redirect(url_for('index'))


@app.errorhandler(404)
def not_found(e):
    """Handle 404 errors"""
    flash('Page not found', 'error')
    return redirect(url_for('index'))


@app.errorhandler(500)
def server_error(e):
    """Handle server errors"""
    flash('An internal error occurred', 'error')
    return redirect(url_for('index'))


if __name__ == '__main__':
    import os
    
    # Get port from environment variable (Hugging Face uses 7860)
    port = int(os.environ.get('PORT', 7860))
    debug_mode = os.environ.get('FLASK_ENV', 'production') != 'production'
    
    print("🚀 Starting CAMS Air Pollution Visualization App")
    print(f"📊 Available at: http://localhost:{port}")
    print("🔧 CDS API Ready:", downloader.is_client_ready())
    
    # Run the Flask app
    app.run(debug=False, host='0.0.0.0', port=port)