"""
Data export for Fractal Correction Engine results.

Exports evolution and navigation results to CSV and JSON formats
for further analysis or publication.
"""

import json
import csv
import os
import numpy as np
from typing import Optional, Dict, Any

from .engine import EvolutionResult, FractalCorrectionEngine
from .navigation import NavigationResult


def export_metrics_csv(
    result: EvolutionResult,
    filepath: str,
) -> str:
    """
    Export time-series metrics to CSV.

    Columns:
      time, fidelity, fidelity_uncorrected, curvature, torsion,
      kappa_error, tau_error, feedback_norm, prediction_fidelity

    Args:
        result: EvolutionResult from engine.evolve()
        filepath: Output file path

    Returns:
        Absolute path of the written file.
    """
    os.makedirs(os.path.dirname(filepath) or '.', exist_ok=True)

    with open(filepath, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow([
            'time', 'fidelity', 'fidelity_uncorrected',
            'curvature', 'torsion',
            'kappa_error', 'tau_error', 'feedback_norm',
            'prediction_fidelity',
        ])
        for i in range(len(result.times)):
            writer.writerow([
                f'{result.times[i]:.10e}',
                f'{result.fidelities[i]:.10f}',
                f'{result.fidelities_uncorrected[i]:.10f}',
                f'{result.curvatures[i]:.10e}',
                f'{result.torsions[i]:.10e}',
                f'{result.kappa_errors[i]:.10e}',
                f'{result.tau_errors[i]:.10e}',
                f'{result.feedback_norms[i]:.10e}',
                f'{result.prediction_fidelities[i]:.10f}',
            ])

    return os.path.abspath(filepath)


def export_summary_json(
    summary: Dict[str, Any],
    filepath: str,
) -> str:
    """
    Export summary statistics to JSON.

    Args:
        summary: Dict from engine.summary(result)
        filepath: Output file path

    Returns:
        Absolute path of the written file.
    """
    os.makedirs(os.path.dirname(filepath) or '.', exist_ok=True)

    # Make everything JSON-serializable
    def _serialize(obj):
        if isinstance(obj, (np.integer,)):
            return int(obj)
        if isinstance(obj, (np.floating, np.float64)):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        if isinstance(obj, dict):
            return {k: _serialize(v) for k, v in obj.items()}
        if isinstance(obj, (list, tuple)):
            return [_serialize(v) for v in obj]
        return obj

    data = _serialize(summary)

    with open(filepath, 'w') as f:
        json.dump(data, f, indent=2)

    return os.path.abspath(filepath)


def export_navigation_csv(
    result: NavigationResult,
    filepath: str,
) -> str:
    """
    Export 11D navigation coordinates to CSV.

    Columns:
      time, coord_0..coord_10, momentum_0..momentum_10, fidelity

    Args:
        result: NavigationResult from NavigationEngine.navigate()
        filepath: Output file path

    Returns:
        Absolute path of the written file.
    """
    os.makedirs(os.path.dirname(filepath) or '.', exist_ok=True)

    with open(filepath, 'w', newline='') as f:
        writer = csv.writer(f)

        header = ['time']
        for d in range(11):
            header.append(f'coord_{result.states[0].dimension_info[d].name}')
        for d in range(11):
            header.append(f'momentum_{result.states[0].dimension_info[d].name}')
        header.append('fidelity')
        writer.writerow(header)

        evo = result.evolution_result
        for i, state in enumerate(result.states):
            row = [f'{state.time:.10e}']
            for d in range(11):
                row.append(f'{state.coordinates[d]:.10e}')
            for d in range(11):
                row.append(f'{state.momenta[d]:.10e}')
            row.append(f'{evo.fidelities[i]:.10f}')
            writer.writerow(row)

    return os.path.abspath(filepath)


def export_full(
    engine: FractalCorrectionEngine,
    result: EvolutionResult,
    output_dir: str,
    nav_result: Optional[NavigationResult] = None,
) -> Dict[str, str]:
    """
    Export all data files to an output directory.

    Creates:
      - metrics.csv: Time-series metrics
      - summary.json: Summary statistics
      - navigation.csv: 11D coordinates (if nav_result provided)

    Returns:
        Dict mapping file type to absolute path.
    """
    os.makedirs(output_dir, exist_ok=True)
    paths = {}

    paths['metrics'] = export_metrics_csv(
        result, os.path.join(output_dir, 'metrics.csv')
    )
    paths['summary'] = export_summary_json(
        engine.summary(result), os.path.join(output_dir, 'summary.json')
    )

    if nav_result is not None:
        paths['navigation'] = export_navigation_csv(
            nav_result, os.path.join(output_dir, 'navigation.csv')
        )

    return paths