#!/usr/bin/env python
"""
Extract Mirror Optical Specifications from FEA Mesh Geometry

This tool analyzes mirror mesh geometry to estimate optical specifications
including focal length, aperture diameter, f-number, and radius of curvature.

Usage:
    # From study directory containing OP2 files
    python -m optimization_engine.tools.extract_mirror_optical_specs .

    # From specific OP2 file
    python -m optimization_engine.tools.extract_mirror_optical_specs path/to/results.op2

    # Save to study README
    python -m optimization_engine.tools.extract_mirror_optical_specs . --update-readme

Output:
    - Console: Optical specifications summary
    - Optional: Updates parent README.md with validated specs

Author: Atomizer Framework
"""

from pathlib import Path
import argparse
import sys

# Add project root to path (tools/ is at project root, so parent is Atomizer/)
sys.path.insert(0, str(Path(__file__).parent.parent))

import numpy as np


def find_op2_file(path: Path) -> Path:
    """Find an OP2 file from path (file or directory)."""
    path = Path(path)

    if path.is_file() and path.suffix.lower() == '.op2':
        return path

    if path.is_dir():
        # Look in common locations
        search_patterns = [
            '**/2_iterations/**/*.op2',
            '**/*.op2',
            '2_iterations/**/*.op2',
            '1_setup/model/*.op2',
        ]

        for pattern in search_patterns:
            op2_files = list(path.glob(pattern))
            if op2_files:
                # Return most recent
                return max(op2_files, key=lambda p: p.stat().st_mtime)

    raise FileNotFoundError(f"No OP2 file found in {path}")


def extract_optical_specs(op2_path: Path, verbose: bool = True) -> dict:
    """
    Extract optical specifications from mirror mesh geometry.

    Args:
        op2_path: Path to OP2 file
        verbose: Print detailed output

    Returns:
        dict with optical specifications
    """
    from optimization_engine.extractors.extract_zernike_opd import (
        ZernikeOPDExtractor,
        estimate_focal_length_from_geometry
    )

    if verbose:
        print(f"Analyzing: {op2_path}")
        print("=" * 60)

    extractor = ZernikeOPDExtractor(op2_path)

    # Get geometry
    geo = extractor.node_geometry
    all_pos = np.array(list(geo.values()))
    x, y, z = all_pos[:, 0], all_pos[:, 1], all_pos[:, 2]

    # Compute radius/diameter
    r = np.sqrt(x**2 + y**2)

    # Estimate focal length
    focal = estimate_focal_length_from_geometry(x, y, z, concave=True)

    # Derived quantities
    diameter = 2 * r.max()
    f_number = focal / diameter
    RoC = 2 * focal  # Radius of curvature
    sag = r.max()**2 / (4 * focal)  # Surface sag at edge
    central_obs = r.min() if r.min() > 1.0 else 0.0  # Central obscuration

    # Parabola fit quality check
    r_sq = x**2 + y**2
    A = np.column_stack([r_sq, np.ones_like(r_sq)])
    coeffs, _, _, _ = np.linalg.lstsq(A, z, rcond=None)
    a, b = coeffs
    z_fit = a * r_sq + b
    rms_error = np.sqrt(np.mean((z - z_fit)**2))

    # Determine fit quality
    if rms_error < 0.1:
        fit_quality = "Excellent"
        fit_note = "Focal length estimate is reliable"
    elif rms_error < 1.0:
        fit_quality = "Good"
        fit_note = "Focal length estimate is reasonably accurate"
    else:
        fit_quality = "Poor"
        fit_note = "Consider using explicit focal length from optical design"

    specs = {
        'aperture_diameter_mm': diameter,
        'aperture_radius_mm': r.max(),
        'focal_length_mm': focal,
        'f_number': f_number,
        'radius_of_curvature_mm': RoC,
        'surface_sag_mm': sag,
        'central_obscuration_mm': central_obs,
        'node_count': len(geo),
        'x_range_mm': (x.min(), x.max()),
        'y_range_mm': (y.min(), y.max()),
        'z_range_mm': (z.min(), z.max()),
        'parabola_fit_rms_mm': rms_error,
        'fit_quality': fit_quality,
        'fit_note': fit_note,
        'source_file': str(op2_path),
    }

    if verbose:
        print()
        print("MIRROR OPTICAL SPECIFICATIONS (from mesh geometry)")
        print("=" * 60)
        print()
        print(f"Aperture Diameter:     {diameter:.1f} mm ({diameter/1000:.3f} m)")
        print(f"Aperture Radius:       {r.max():.1f} mm")
        if central_obs > 0:
            print(f"Central Obscuration:   {central_obs:.1f} mm")
        print()
        print(f"Estimated Focal Length: {focal:.1f} mm ({focal/1000:.3f} m)")
        print(f"Radius of Curvature:    {RoC:.1f} mm ({RoC/1000:.3f} m)")
        print(f"f-number (f/D):         f/{f_number:.2f}")
        print()
        print(f"Surface Sag at Edge:    {sag:.2f} mm")
        print()
        print("--- Mesh Statistics ---")
        print(f"Node count:    {len(geo)}")
        print(f"X range:       {x.min():.1f} to {x.max():.1f} mm")
        print(f"Y range:       {y.min():.1f} to {y.max():.1f} mm")
        print(f"Z range:       {z.min():.2f} to {z.max():.2f} mm")
        print()
        print("--- Parabola Fit Quality ---")
        print(f"RMS fit residual: {rms_error:.4f} mm ({rms_error*1000:.2f} µm)")
        print(f"Quality: {fit_quality} - {fit_note}")
        print()
        print("=" * 60)

    return specs


def generate_readme_section(specs: dict) -> str:
    """Generate markdown section for README."""
    return f"""## 2. Optical Prescription

> **Source**: Estimated from mesh geometry. Validate against optical design.

| Parameter | Value | Units | Status |
|-----------|-------|-------|--------|
| Aperture Diameter | {specs['aperture_diameter_mm']:.1f} | mm | Estimated |
| Focal Length | {specs['focal_length_mm']:.1f} | mm | Estimated |
| f-number | f/{specs['f_number']:.2f} | - | Computed |
| Radius of Curvature | {specs['radius_of_curvature_mm']:.1f} | mm | Computed (2×f) |
| Central Obscuration | {specs['central_obscuration_mm']:.1f} | mm | From mesh |
| Surface Type | Parabola | - | Assumed |

**Fit Quality**: {specs['fit_quality']} ({specs['fit_note']})

### 2.1 Usage in OPD Extractor

For rigorous WFE analysis, use explicit focal length:

```python
from optimization_engine.extractors import ZernikeOPDExtractor

extractor = ZernikeOPDExtractor(
    op2_file,
    focal_length={specs['focal_length_mm']:.1f},  # mm - validate against design
    concave=True
)
```
"""


def update_readme(study_dir: Path, specs: dict):
    """Update parent README with optical specs."""
    readme_path = study_dir / 'README.md'

    if not readme_path.exists():
        print(f"README.md not found at {readme_path}")
        return False

    content = readme_path.read_text(encoding='utf-8')

    # Find and replace optical prescription section
    new_section = generate_readme_section(specs)

    # Look for existing section
    import re
    pattern = r'## 2\. Optical Prescription.*?(?=## 3\.|$)'

    if re.search(pattern, content, re.DOTALL):
        content = re.sub(pattern, new_section + '\n---\n\n', content, flags=re.DOTALL)
        print(f"Updated optical prescription in {readme_path}")
    else:
        print(f"Could not find '## 2. Optical Prescription' section in {readme_path}")
        print("Please add manually or check section numbering.")
        return False

    readme_path.write_text(content, encoding='utf-8')
    return True


def main():
    parser = argparse.ArgumentParser(
        description='Extract mirror optical specifications from FEA mesh',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    # Analyze current study directory
    python -m optimization_engine.tools.extract_mirror_optical_specs .

    # Analyze specific OP2 file
    python -m optimization_engine.tools.extract_mirror_optical_specs results.op2

    # Update parent README with specs
    python -m optimization_engine.tools.extract_mirror_optical_specs . --update-readme
        """
    )

    parser.add_argument('path', type=str,
                       help='Path to OP2 file or study directory')
    parser.add_argument('--update-readme', action='store_true',
                       help='Update parent README.md with optical specs')
    parser.add_argument('--quiet', '-q', action='store_true',
                       help='Suppress detailed output')
    parser.add_argument('--json', action='store_true',
                       help='Output specs as JSON')

    args = parser.parse_args()

    try:
        path = Path(args.path).resolve()
        op2_path = find_op2_file(path)

        specs = extract_optical_specs(op2_path, verbose=not args.quiet and not args.json)

        if args.json:
            import json
            print(json.dumps(specs, indent=2, default=str))

        if args.update_readme:
            # Find study root (parent of geometry type folder)
            study_dir = path if path.is_dir() else path.parent
            # Go up to geometry type level
            while study_dir.name not in ['studies', ''] and not (study_dir / 'README.md').exists():
                if (study_dir.parent / 'README.md').exists():
                    study_dir = study_dir.parent
                    break
                study_dir = study_dir.parent

            if (study_dir / 'README.md').exists():
                update_readme(study_dir, specs)
            else:
                print(f"Could not find parent README.md to update")

    except FileNotFoundError as e:
        print(f"Error: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()