import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Union
from dataclasses import dataclass
import matplotlib.pyplot as plt
import seaborn as sns

@dataclass
class ConductorParams:
    """Data class for conductor electrical parameters"""
    resistance: float  # Ω/mile
    gmr: float        # feet (Geometric Mean Radius)

@dataclass
class Coordinate:
    """Data class for conductor coordinates"""
    x: float  # feet
    y: float  # feet

class PowerSystemImpedanceCalculator:
    """
    Advanced impedance calculator for 5-wire power systems using modified Carson's equations
    and Kron reduction technique.
    
    The calculator implements industry-standard methods for computing equivalent impedance
    matrices in multi-conductor transmission and distribution lines.
    """

    # Carson's equation constants for 60 Hz, 100 Ω⋅m earth resistivity
    CARSON_REAL_CONSTANT = 0.09530  # Ω/mile
    CARSON_IMAG_COEFFICIENT = 0.12134  # Ω/mile
    CARSON_IMAG_CONSTANT = 7.93402  # dimensionless

    def __init__(self):
        """Initialize the calculator with default conductor labels"""
        self.conductor_labels = ['a', 'b', 'c', 'n', 'pe']
        self.phase_labels = ['a', 'b', 'c']
        
    def calculate_distance_from_coordinates(self, coord1: Coordinate, coord2: Coordinate) -> float:
        """
        Calculate Euclidean distance between two conductor coordinates.
        
        Args:
            coord1: First conductor coordinate
            coord2: Second conductor coordinate
            
        Returns:
            Distance in feet
        """
        dx = coord1.x - coord2.x
        dy = coord1.y - coord2.y
        return np.sqrt(dx**2 + dy**2)

    def calculate_all_distances_from_coordinates(self, coordinates: Dict[str, Coordinate]) -> Dict[str, float]:
        """
        Calculate all pairwise distances from conductor coordinates.
        
        Args:
            coordinates: Dictionary mapping conductor labels to coordinates
            
        Returns:
            Dictionary of pairwise distances with keys like 'ab', 'ac', etc.
        """
        distances = {}
        conductors = self.conductor_labels
        
        # Generate all unique pairs
        for i, cond1 in enumerate(conductors):
            for j, cond2 in enumerate(conductors[i+1:], i+1):
                key = f"{cond1}{cond2}"
                distance = self.calculate_distance_from_coordinates(
                    coordinates[cond1], coordinates[cond2]
                )
                distances[key] = distance
                
        return distances

    def calculate_primitive_impedance_matrix(self, 
                                         distances: Dict[str, float], 
                                         conductor_params: Dict[str, ConductorParams]) -> np.ndarray:
        """
        Calculate the 5×5 primitive impedance matrix using modified Carson's equations.
        
        The primitive matrix includes all conductors (phases, neutral, PE) before reduction.
        Self-impedances account for conductor resistance and earth return effects.
        Mutual impedances account for electromagnetic coupling and earth return effects.
        
        Args:
            distances: Dictionary of pairwise distances between conductors
            conductor_params: Dictionary of conductor electrical parameters
            
        Returns:
            5×5 complex impedance matrix [Ω/mile]
        """
        n_conductors = len(self.conductor_labels)
        matrix = np.zeros((n_conductors, n_conductors), dtype=complex)
        
        # Calculate self-impedances (diagonal elements)
        for i, conductor in enumerate(self.conductor_labels):
            params = conductor_params[conductor]
            
            # Modified Carson's equation for self-impedance
            real_part = params.resistance + self.CARSON_REAL_CONSTANT
            imag_part = self.CARSON_IMAG_COEFFICIENT * (
                np.log(1.0 / params.gmr) + self.CARSON_IMAG_CONSTANT
            )
            
            matrix[i, i] = complex(real_part, imag_part)
        
        # Calculate mutual impedances (off-diagonal elements)
        # Mapping from matrix indices to distance dictionary keys
        distance_map = {
            (0, 1): 'ab', (0, 2): 'ac', (0, 3): 'an', (0, 4): 'ape',
            (1, 2): 'bc', (1, 3): 'bn', (1, 4): 'bpe',
            (2, 3): 'cn', (2, 4): 'cpe',
            (3, 4): 'npe'
        }
        
        for (i, j), distance_key in distance_map.items():
            distance = distances[distance_key]
            
            # Modified Carson's equation for mutual impedance
            real_part = self.CARSON_REAL_CONSTANT
            imag_part = self.CARSON_IMAG_COEFFICIENT * (
                np.log(1.0 / distance) + self.CARSON_IMAG_CONSTANT
            )
            
            mutual_impedance = complex(real_part, imag_part)
            matrix[i, j] = mutual_impedance
            matrix[j, i] = mutual_impedance  # Symmetry
            
        return matrix

    def apply_kron_reduction(self, primitive_matrix: np.ndarray) -> np.ndarray:
        """
        Apply Kron reduction to eliminate neutral and PE conductors from the impedance matrix.
        
        Kron reduction preserves the electrical behavior of the phase conductors while
        eliminating the neutral and protective earth conductors from the analysis.
        The reduction formula is: Z_abc = Z_pp - Z_pq @ Z_qq^(-1) @ Z_qp
        
        Args:
            primitive_matrix: 5×5 primitive impedance matrix
            
        Returns:
            3×3 reduced impedance matrix for phase conductors only [Ω/mile]
        """
        # Extract sub-matrices for Kron reduction
        # Z_pp: phase-to-phase impedances (3×3) - indices 0,1,2 (a,b,c)
        # Z_qq: neutral/PE impedances (2×2) - indices 3,4 (n,pe)  
        # Z_pq: phase-to-neutral/PE coupling (3×2)
        # Z_qp: neutral/PE-to-phase coupling (2×3) - transpose of Z_pq
        
        Z_pp = primitive_matrix[0:3, 0:3]  # Phase conductors
        Z_qq = primitive_matrix[3:5, 3:5]  # Neutral and PE
        Z_pq = primitive_matrix[0:3, 3:5]  # Phase to neutral/PE coupling
        Z_qp = primitive_matrix[3:5, 0:3]  # Neutral/PE to phase coupling
        
        # Calculate Z_qq inverse using robust numerical methods
        try:
            Z_qq_inv = np.linalg.inv(Z_qq)
        except np.linalg.LinAlgError:
            # Fallback to pseudo-inverse if matrix is singular
            Z_qq_inv = np.linalg.pinv(Z_qq)
            print("Warning: Z_qq matrix is singular, using pseudo-inverse")
        
        # Apply Kron reduction formula
        # Z_abc = Z_pp - Z_pq @ Z_qq^(-1) @ Z_qp
        correction_term = Z_pq @ Z_qq_inv @ Z_qp
        reduced_matrix = Z_pp - correction_term
        
        return reduced_matrix

    def calculate_impedance_from_distances(self, 
                                       distances: Dict[str, float],
                                       conductor_params: Dict[str, ConductorParams]) -> Tuple[np.ndarray, np.ndarray]:
        """
        Complete impedance calculation from conductor distances.
        
        Args:
            distances: Dictionary of pairwise conductor distances
            conductor_params: Dictionary of conductor electrical parameters
            
        Returns:
            Tuple of (primitive_matrix, reduced_matrix)
        """
        primitive_matrix = self.calculate_primitive_impedance_matrix(distances, conductor_params)
        reduced_matrix = self.apply_kron_reduction(primitive_matrix)
        
        return primitive_matrix, reduced_matrix

    def calculate_impedance_from_coordinates(self, 
                                         coordinates: Dict[str, Coordinate],
                                         conductor_params: Dict[str, ConductorParams]) -> Tuple[np.ndarray, np.ndarray, Dict[str, float]]:
        """
        Complete impedance calculation from conductor coordinates.
        
        Args:
            coordinates: Dictionary mapping conductor labels to coordinates
            conductor_params: Dictionary of conductor electrical parameters
            
        Returns:
            Tuple of (primitive_matrix, reduced_matrix, calculated_distances)
        """
        distances = self.calculate_all_distances_from_coordinates(coordinates)
        primitive_matrix, reduced_matrix = self.calculate_impedance_from_distances(distances, conductor_params)
        
        return primitive_matrix, reduced_matrix, distances

    def format_complex_matrix(self, matrix: np.ndarray, precision: int = 6) -> List[List[str]]:
        """
        Format complex matrix for readable display.
        
        Args:
            matrix: Complex numpy array
            precision: Number of decimal places
            
        Returns:
            List of lists containing formatted complex number strings
        """
        formatted = []
        for row in matrix:
            formatted_row = []
            for element in row:
                real = f"{element.real:.{precision}f}"
                imag = f"{element.imag:.{precision}f}"
                sign = "+" if element.imag >= 0 else ""
                formatted_row.append(f"{real} {sign} j{imag}")
            formatted.append(formatted_row)
        
        return formatted

    def create_impedance_dataframe(self, matrix: np.ndarray, labels: List[str]) -> pd.DataFrame:
        """
        Create a pandas DataFrame from impedance matrix for easier analysis.
        
        Args:
            matrix: Complex impedance matrix
            labels: Row/column labels
            
        Returns:
            DataFrame with complex impedance values
        """
        return pd.DataFrame(matrix, index=labels, columns=labels)

    def analyze_matrix_properties(self, matrix: np.ndarray) -> Dict[str, Union[float, bool]]:
        """
        Analyze impedance matrix properties for engineering insights.
        
        Args:
            matrix: Complex impedance matrix
            
        Returns:
            Dictionary containing matrix analysis results
        """
        properties = {}
        
        # Basic properties
        properties['condition_number'] = np.linalg.cond(matrix)
        properties['determinant'] = np.linalg.det(matrix)
        properties['is_symmetric'] = np.allclose(matrix, matrix.T, rtol=1e-10)
        properties['is_positive_definite'] = np.all(np.linalg.eigvals(matrix.real) > 0)
        
        # Eigenvalue analysis
        eigenvalues = np.linalg.eigvals(matrix)
        properties['min_eigenvalue_real'] = np.min(eigenvalues.real)
        properties['max_eigenvalue_real'] = np.max(eigenvalues.real)
        properties['eigenvalue_spread'] = properties['max_eigenvalue_real'] / properties['min_eigenvalue_real']
        
        # Impedance magnitude analysis
        magnitudes = np.abs(matrix)
        properties['min_impedance_magnitude'] = np.min(magnitudes)
        properties['max_impedance_magnitude'] = np.max(magnitudes)
        properties['avg_self_impedance_magnitude'] = np.mean(np.abs(np.diag(matrix)))
        
        return properties

    def format_results_for_display(self, primitive_matrix: np.ndarray, 
                                 reduced_matrix: np.ndarray, 
                                 distances: Dict[str, float]) -> str:
        """
        Format calculation results for display in Gradio interface.
        
        Args:
            primitive_matrix: 5×5 primitive impedance matrix
            reduced_matrix: 3×3 reduced impedance matrix
            distances: Dictionary of conductor distances
            
        Returns:
            Formatted string with results
        """
        result = "=== 3-Phase Impedance Calculation Results ===\n\n"
        
        # Display distances
        result += "Conductor Distances:\n"
        for pair, distance in distances.items():
            conductor1, conductor2 = pair[0].upper(), pair[1].upper()
            result += f"  {conductor1}-{conductor2}: {distance:.3f} ft\n"
        
        result += "\nPrimitive Impedance Matrix (5×5) [Ω/mile]:\n"
        result += "     "
        for label in self.conductor_labels:
            result += f"{label.upper():>15s}"
        result += "\n"
        
        for i, label in enumerate(self.conductor_labels):
            result += f"{label.upper():>3s}: "
            for j in range(5):
                z = primitive_matrix[i, j]
                result += f"{z.real:>7.3f}{z.imag:>+7.3f}j "
            result += "\n"
        
        result += "\nReduced Phase Impedance Matrix (3×3) [Ω/mile]:\n"
        result += "     "
        for label in self.phase_labels:
            result += f"{label.upper():>15s}"
        result += "\n"
        
        for i, label in enumerate(self.phase_labels):
            result += f"{label.upper():>3s}: "
            for j in range(3):
                z = reduced_matrix[i, j]
                result += f"{z.real:>7.3f}{z.imag:>+7.3f}j "
            result += "\n"
        
        # Add matrix analysis
        properties = self.analyze_matrix_properties(reduced_matrix)
        result += f"\nMatrix Analysis:\n"
        result += f"  Condition Number: {properties['condition_number']:.2e}\n"
        result += f"  Is Symmetric: {properties['is_symmetric']}\n"
        result += f"  Average Self-Impedance Magnitude: {properties['avg_self_impedance_magnitude']:.5f} Ω/mile\n"
        
        return result

def create_test_data():
    """Create sample data for testing the calculator"""
    
    # Sample conductor parameters (typical values)
    conductor_params = {
        'a': ConductorParams(resistance=0.055, gmr=0.038),    # 4/0 ACSR
        'b': ConductorParams(resistance=0.055, gmr=0.038),    # 4/0 ACSR  
        'c': ConductorParams(resistance=0.055, gmr=0.038),    # 4/0 ACSR
        'n': ConductorParams(resistance=8.0, gmr=0.012),      # #6 ACSR
        'pe': ConductorParams(resistance=8.0, gmr=0.012)      # #6 ACSR
    }

    # Sample coordinates (typical distribution line configuration)
    coordinates = {
        'a': Coordinate(x=0, y=42),
        'b': Coordinate(x=23.5, y=42), 
        'c': Coordinate(x=47, y=42),
        'n': Coordinate(x=10, y=74),
        'pe': Coordinate(x=37, y=72)
    }

    return conductor_params, coordinates