#!/usr/bin/env python3
"""
CDZ mapping automation script.
Maps Creative Dissolution Zones from survey data.
"""

import argparse
import sys
import numpy as np
import pandas as pd
import json
from pathlib import Path

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

try:
    from cavora.cdz import CDZDetector, CDZMapper, CDZClassifier
    from cavora.utils.logger import setup_logger
except ImportError:
    print("Error: CAVORA package not found")
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description="Map Creative Dissolution Zones")
    parser.add_argument("--input", "-i", required=True, help="Input survey data file")
    parser.add_argument("--output", "-o", default="cdz_map.geojson", help="Output GeoJSON file")
    parser.add_argument("--threshold", "-t", type=float, default=1.5, help="Divergence threshold")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    
    args = parser.parse_args()
    
    # Setup logging
    logger = setup_logger("cdz_mapping", 
                         log_level="DEBUG" if args.verbose else "INFO")
    
    print("\n" + "="*60)
    print("CAVORA - CDZ Mapping")
    print("="*60)
    
    # Load data
    print(f"\nLoading data from: {args.input}")
    try:
        if args.input.endswith('.csv'):
            df = pd.read_csv(args.input)
        elif args.input.endswith('.json'):
            with open(args.input, 'r') as f:
                data = json.load(f)
                df = pd.DataFrame(data)
        else:
            print(f"Unsupported file format: {args.input}")
            sys.exit(1)
        
        print(f"Loaded {len(df)} data points")
    except Exception as e:
        print(f"Error loading data: {e}")
        sys.exit(1)
    
    # Check required columns
    required = ['passage_id', 'time', 'V']
    missing = [col for col in required if col not in df.columns]
    if missing:
        print(f"Missing required columns: {missing}")
        sys.exit(1)
    
    # Initialize detectors
    detector = CDZDetector(divergence_threshold=args.threshold)
    mapper = CDZMapper()
    classifier = CDZClassifier()
    
    # Process each passage
    passages = df['passage_id'].unique()
    all_cdzs = []
    
    print(f"\nAnalyzing {len(passages)} passages for CDZs...")
    
    for passage in passages:
        print(f"  Processing: {passage}")
        
        # Get data for this passage
        passage_data = df[df['passage_id'] == passage].sort_values('time')
        
        if len(passage_data) < 5:
            continue
        
        # Extract time series
        t = passage_data['time'].values
        V = passage_data['V'].values
        L = passage_data['L'].values if 'L' in passage_data.columns else None
        F = passage_data['F'].values if 'F' in passage_data.columns else None
        CO2 = passage_data['CO2'].values if 'CO2' in passage_data.columns else None
        
        # Detect CDZs
        cdzs = detector.detect_from_trajectory(
            t, V, L, F, CO2, location=passage
        )
        
        if cdzs:
            print(f"    Found {len(cdzs)} CDZs")
            
            for cdz in cdzs:
                # Get coordinates if available
                x = passage_data['x'].values[0] if 'x' in passage_data.columns else 0
                y = passage_data['y'].values[0] if 'y' in passage_data.columns else 0
                z = passage_data['z'].values[0] if 'z' in passage_data.columns else 0
                
                # Add to mapper
                mapper.add_cdz_location(
                    x=x, y=y, z=z,
                    zone_type=cdz['transition_type'],
                    size=cdz['duration'] * 100,  # Convert time to approximate size
                    intensity=cdz['confidence'],
                    passage_id=passage
                )
                
                all_cdzs.append(cdz)
    
    # Generate map
    print(f"\nGenerating CDZ map...")
    
    # Create GeoJSON
    features = []
    for i, cdz in enumerate(all_cdzs):
        feature = {
            "type": "Feature",
            "geometry": {
                "type": "Point",
                "coordinates": [mapper.locations[i].x, mapper.locations[i].y, mapper.locations[i].z]
            },
            "properties": {
                "id": f"CDZ_{i+1:03d}",
                "passage": cdz['location'],
                "type": cdz['transition_type'],
                "start_time": cdz['start_time'],
                "end_time": cdz['end_time'],
                "duration": cdz['duration'],
                "confidence": cdz['confidence']
            }
        }
        features.append(feature)
    
    geojson = {
        "type": "FeatureCollection",
        "features": features
    }
    
    with open(args.output, 'w') as f:
        json.dump(geojson, f, indent=2)
    
    print(f"Map saved to: {args.output}")
    print(f"\nCDZ Statistics:")
    print(f"  Total CDZs detected: {len(all_cdzs)}")
    
    if all_cdzs:
        # Type distribution
        types = {}
        for cdz in all_cdzs:
            t = cdz['transition_type']
            types[t] = types.get(t, 0) + 1
        
        print("\n  Type distribution:")
        for t, count in types.items():
            print(f"    {t}: {count} ({count/len(all_cdzs)*100:.1f}%)")
    
    print("\n" + "="*60)


if __name__ == "__main__":
    main()
