import numpy as np
import biotite.structure as struc
import csv
import json
import logging
from typing import List, Dict, Union, Any

logger = logging.getLogger(__name__)

def calculate_torsion_angles(atom_array: struc.AtomArray) -> List[Dict[str, Any]]:
    """
    Calculate backbone torsion angles (phi, psi, omega) for a given structure.
    
    Args:
        atom_array: Biotite AtomArray containing the protein structure.
        
    Returns:
        List of dictionaries, one per residue, containing:
        - residue: 3-letter code
        - res_id: Residue number
        - phi: Phi angle in degrees (or None)
        - psi: Psi angle in degrees (or None)
        - omega: Omega angle in degrees (or None)
    """
    logger.info("Calculating backbone torsion angles...")
    
    # Calculate angles using Biotite
    # dihedral_backbone returns (phi, psi, omega) in RADIANS for each residue.
    # Note: biotite.structure.dihedral_backbone calculates phi, psi, omega for each residue.
    # The first residue's phi is usually NaN, last residue's psi is NaN.
    
    phi, psi, omega = struc.dihedral_backbone(atom_array)
    
    # Convert to degrees
    phi_deg = np.degrees(phi)
    psi_deg = np.degrees(psi)
    omega_deg = np.degrees(omega)
    
    # Get residue identifiers
    # We iterate through indices corresponding to the backbone results
    # Ideally should use residue iterator to be safe, but backbone arrays map to residues.
    
    # Get unique residue IDs to map back
    residue_starts = struc.get_residue_starts(atom_array)
    
    results = []
    
    for i, start_idx in enumerate(residue_starts):
        res_name = atom_array.res_name[start_idx]
        res_id = int(atom_array.res_id[start_idx])
        
        # Handle NaN values (convert to None for JSON compliance)
        p = phi_deg[i]
        ps = psi_deg[i]
        o = omega_deg[i]
        
        if np.isnan(p): p = None
        if np.isnan(ps): ps = None
        if np.isnan(o): o = None
        
        entry = {
            'residue': res_name,
            'res_id': res_id,
            'phi': p,
            'psi': ps,
            'omega': o
        }
        results.append(entry)
        
    logger.info(f"Calculated angles for {len(results)} residues.")
    return results

def export_torsion_angles(data: List[Dict[str, Any]], output_file: str, fmt: str = "csv"):
    """
    Export calculated torsion angles to a file.
    
    Args:
        data: List of angle dictionaries (from calculate_torsion_angles).
        output_file: Path to output file.
        fmt: Format 'csv' or 'json'.
    """
    logger.info(f"Exporting torsion angles to {output_file} ({fmt})...")
    
    if fmt.lower() == "json":
        with open(output_file, 'w') as f:
            json.dump(data, f, indent=2)
            
    elif fmt.lower() == "csv":
        fieldnames = ['residue', 'res_id', 'phi', 'psi', 'omega']
        with open(output_file, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            for row in data:
                # Ensure None becomes empty string or specific marker if needed
                # csv module handles None as empty string by default? No, usually expects string.
                # Actually DictWriter handles it if we don't convert.
                # But let's be explicit solely for robustness if needed.
                writer.writerow(row)
    else:
        raise ValueError(f"Unknown format: {fmt}")
        
    logger.info("Export complete.")
