"""
Plotting utilities for spin wave spectra.
"""

import numpy as np
from typing import Dict, Optional, Tuple, List
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize, LogNorm


def sw_plotspec(spec: Dict, mode: str = 'color',
                axLim: List[float] = None,
                dE: float = 0.1,
                colormap: str = 'viridis',
                colorbar: bool = True,
                dashed: bool = False,
                imag: bool = False,
                log: bool = False,
                qlabel: List[str] = None,
                title: str = None,
                ax: plt.Axes = None,
                **kwargs) -> plt.Axes:
    """
    Plot spin wave spectrum.
    
    Parameters
    ----------
    spec : dict
        Output from SpinW.spinwave() or sw_egrid()
    mode : str
        Plot mode:
        - 'disp': Plot dispersion curves only
        - 'color': Color intensity map
        - 'line': Line plot of intensity
        - 'waterfall': Waterfall plot
    axLim : list
        Axis limits [Emin, Emax] or [xmin, xmax, Emin, Emax]
    dE : float
        Energy resolution for convolution (if not done)
    colormap : str
        Matplotlib colormap name
    colorbar : bool
        Show colorbar for color mode
    dashed : bool
        Use dashed lines for negative energy modes
    imag : bool
        Plot imaginary part of frequencies
    log : bool
        Use logarithmic color scale
    qlabel : list
        Labels for Q-point ticks
    title : str
        Plot title
    ax : Axes
        Matplotlib axes to plot on
        
    Returns
    -------
    ax : Axes
        Matplotlib axes object
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 6))
    
    omega = spec['omega']  # (n_q, n_modes)
    hkl = spec['hkl']
    n_q = len(hkl)
    
    # Create x-axis (cumulative distance in Q-space)
    x = np.zeros(n_q)
    for i in range(1, n_q):
        x[i] = x[i-1] + np.linalg.norm(hkl[i] - hkl[i-1])
    
    if mode == 'disp' or mode == '1':
        # Plot dispersion curves
        omega_real = np.real(omega)
        omega_imag = np.imag(omega)
        
        for imode in range(omega.shape[1]):
            # Positive energy branch
            ax.plot(x, omega_real[:, imode], 'b-', linewidth=1.5)
            
            if imag and np.any(np.abs(omega_imag[:, imode]) > 1e-6):
                ax.plot(x, omega_imag[:, imode], 'r.', markersize=3)
            
            if dashed:
                # Negative energy (hole) branch
                ax.plot(x, -omega_real[:, imode], 'b--', linewidth=1.0)
        
        ax.set_ylabel('Energy (meV)')
        ax.axhline(y=0, color='k', linewidth=0.5)
    
    elif mode == 'color' or mode == '3':
        # Color intensity map
        if 'swConv' not in spec:
            from .egrid import sw_egrid
            spec = sw_egrid(spec, dE=dE)
        
        swConv = spec['swConv']  # (n_E, n_q)
        Evect = spec['Evect']
        
        # Create meshgrid
        X, Y = np.meshgrid(x, Evect)
        
        # Normalize
        norm = LogNorm(vmin=1e-3, vmax=np.max(swConv)) if log else Normalize()
        
        im = ax.pcolormesh(X, Y, swConv, cmap=colormap, norm=norm, shading='auto')
        
        if colorbar:
            plt.colorbar(im, ax=ax, label='Intensity (arb. units)')
        
        ax.set_ylabel('Energy (meV)')
    
    elif mode == 'line' or mode == '2':
        # Line plot of intensity
        if 'swInt' not in spec:
            from .egrid import sw_egrid
            spec = sw_egrid(spec, dE=dE)
        
        swInt = np.real(spec['swInt'])
        omega_real = np.real(omega)
        
        for imode in range(omega.shape[1]):
            sizes = swInt[:, imode] / np.max(swInt[:, imode]) * 50
            ax.scatter(x, omega_real[:, imode], s=sizes, alpha=0.5)
        
        ax.set_ylabel('Energy (meV)')
    
    else:
        raise ValueError(f"Unknown mode: {mode}")
    
    # Set axis limits
    if axLim is not None:
        if len(axLim) == 2:
            ax.set_ylim(axLim)
        elif len(axLim) == 4:
            ax.set_xlim(axLim[:2])
            ax.set_ylim(axLim[2:])
    
    # Add Q-labels
    if qlabel is not None:
        # Find tick positions (endpoints and segment changes)
        tick_pos = [x[0], x[-1]]
        tick_labels = [qlabel[0], qlabel[-1]] if len(qlabel) >= 2 else qlabel
        ax.set_xticks(tick_pos)
        ax.set_xticklabels(tick_labels)
    else:
        ax.set_xlabel('Q (r.l.u.)')
    
    if title:
        ax.set_title(title)
    
    return ax


def plot_structure(sw, ax: plt.Axes = None, range: List = None,
                   show_bonds: bool = True, show_spins: bool = True,
                   **kwargs) -> plt.Axes:
    """
    Plot the magnetic structure.
    
    Parameters
    ----------
    sw : SpinW
        SpinW object
    ax : Axes
        Matplotlib axes (3D projection)
    range : list
        Plot range [[xmin, xmax], [ymin, ymax], [zmin, zmax]]
    show_bonds : bool
        Show exchange bonds
    show_spins : bool
        Show spin vectors
        
    Returns
    -------
    ax : Axes
        Matplotlib axes object
    """
    from mpl_toolkits.mplot3d import Axes3D
    
    if ax is None:
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(111, projection='3d')
    
    if range is None:
        range = [[0, 1], [0, 1], [0, 1]]
    
    # Plot atoms
    for atom in sw._atoms:
        pos = sw._lattice.frac_to_cart(atom.position)
        if atom.spin > 0:
            ax.scatter(*pos, c='red', s=100, marker='o', label=atom.label)
        else:
            ax.scatter(*pos, c='gray', s=50, marker='o')
    
    # Plot spins
    if show_spins and sw._magstr.S is not None:
        mag_atoms = [a for a in sw._atoms if a.spin > 0]
        for i, atom in enumerate(mag_atoms):
            pos = sw._lattice.frac_to_cart(atom.position)
            spin = sw._magstr.S[i]
            ax.quiver(*pos, *spin, color='blue', arrow_length_ratio=0.3)
    
    # Plot bonds
    if show_bonds:
        for bond in sw._coupling.bonds[:50]:  # Limit for visibility
            if bond.exchange is not None:
                pos1 = sw._lattice.frac_to_cart(sw._atoms[bond.atom1].position)
                pos2 = sw._lattice.frac_to_cart(
                    sw._atoms[bond.atom2].position + bond.dl
                )
                ax.plot([pos1[0], pos2[0]], [pos1[1], pos2[1]], 
                       [pos1[2], pos2[2]], 'g-', alpha=0.5)
    
    ax.set_xlabel('x (Å)')
    ax.set_ylabel('y (Å)')
    ax.set_zlabel('z (Å)')
    
    return ax


def sw_plotspec_multi(specs: List[Dict], labels: List[str] = None,
                      mode: str = 'disp', **kwargs) -> plt.Axes:
    """
    Compare multiple spectra on same axes.
    
    Parameters
    ----------
    specs : list
        List of spectrum dictionaries
    labels : list
        Legend labels for each spectrum
    mode : str
        Plot mode (same as sw_plotspec)
        
    Returns
    -------
    ax : Axes
        Matplotlib axes object
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    
    colors = plt.cm.tab10(np.linspace(0, 1, len(specs)))
    
    for i, spec in enumerate(specs):
        omega = np.real(spec['omega'])
        hkl = spec['hkl']
        n_q = len(hkl)
        
        x = np.zeros(n_q)
        for j in range(1, n_q):
            x[j] = x[j-1] + np.linalg.norm(hkl[j] - hkl[j-1])
        
        for imode in range(omega.shape[1]):
            label = labels[i] if (labels and imode == 0) else None
            ax.plot(x, omega[:, imode], color=colors[i], linewidth=1.5, label=label)
    
    if labels:
        ax.legend()
    
    ax.set_xlabel('Q (r.l.u.)')
    ax.set_ylabel('Energy (meV)')
    
    return ax
