#!/usr/bin/env python3
"""
Flood pulse prediction script.
Forecasts flood events from precipitation data.
"""

import argparse
import sys
import json
from datetime import datetime

sys.path.insert(0, '.')

try:
    from cavora.prediction import FloodForecasting
    from cavora.utils.logger import setup_logger
except ImportError:
    print("Error: CAVORA package not found")
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description="Predict flood pulses from precipitation")
    parser.add_argument("--precip", "-p", required=True, help="Precipitation data (JSON file)")
    parser.add_argument("--current-q", type=float, default=0.15, help="Current discharge (m³/s)")
    parser.add_argument("--current-level", type=float, default=0.45, help="Current water level (m)")
    parser.add_argument("--soil-sat", type=float, default=0.65, help="Soil saturation (0-1)")
    parser.add_argument("--output", "-o", help="Output file (JSON)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    
    args = parser.parse_args()
    
    # Setup logging
    logger = setup_logger("flood_prediction", 
                         log_level="DEBUG" if args.verbose else "INFO")
    
    print("\n" + "="*60)
    print("CAVORA - Flood Pulse Prediction")
    print("="*60)
    
    # Load precipitation data
    print(f"\nLoading precipitation data from: {args.precip}")
    try:
        with open(args.precip, 'r') as f:
            precip_data = json.load(f)
        
        # Convert to list format if needed
        if isinstance(precip_data, dict):
            if 'events' in precip_data:
                precip_events = precip_data['events']
            else:
                # Assume it's a single event
                precip_events = [precip_data]
        else:
            precip_events = precip_data
        
        print(f"Loaded {len(precip_events)} precipitation events")
        for i, event in enumerate(precip_events):
            print(f"  Event {i+1}: {event.get('amount', 0)}mm at t={event.get('time', 0)}h")
            
    except Exception as e:
        print(f"Error loading precipitation data: {e}")
        sys.exit(1)
    
    # Initialize flood forecaster
    ff = FloodForecasting()
    
    # Generate forecast
    print(f"\nGenerating flood forecast...")
    forecast = ff.forecast_from_precipitation(
        precip_events,
        current_discharge=args.current_q,
        current_level=args.current_level,
        soil_saturation=args.soil_sat
    )
    
    # Display results
    print(f"\n{'='*60}")
    print(f"FLOOD FORECAST")
    print(f"{'='*60}")
    print(f"\nSeverity: {forecast.severity}")
    print(f"Peak Discharge: {forecast.peak_discharge} m³/s")
    print(f"Peak Water Level: {forecast.peak_level} m")
    print(f"Time to Peak: {forecast.time_to_peak} hours")
    print(f"Duration: {forecast.duration} hours")
    print(f"Probability: {forecast.probability:.0%}")
    
    # Evacuation timing
    print(f"\n{'='*60}")
    print(f"EVACUATION ASSESSMENT")
    print(f"{'='*60}")
    
    timing = ff.get_evacuation_timing(forecast, exit_time=2.0)
    print(f"\nAction: {timing['action']}")
    print(f"Time Available: {timing['time_available_hours']} hours")
    print(f"Evacuate Now: {'YES' if timing['evacuate_now'] else 'NO'}")
    
    # Save results
    if args.output:
        output_data = {
            'timestamp': datetime.now().isoformat(),
            'input': {
                'precipitation_events': precip_events,
                'current_discharge': args.current_q,
                'current_level': args.current_level,
                'soil_saturation': args.soil_sat
            },
            'forecast': {
                'severity': forecast.severity,
                'peak_discharge': forecast.peak_discharge,
                'peak_level': forecast.peak_level,
                'time_to_peak': forecast.time_to_peak,
                'duration': forecast.duration,
                'probability': forecast.probability
            },
            'evacuation': timing
        }
        
        with open(args.output, 'w') as f:
            json.dump(output_data, f, indent=2)
        
        print(f"\nForecast saved to: {args.output}")
    
    print("\n" + "="*60)


if __name__ == "__main__":
    main()
