"""
Visualization utilities for macrolactone analysis.
"""
from rdkit import Chem
from rdkit.Chem import Draw
from typing import Dict, List
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np


def draw_molecule_with_numbering(mol: Chem.Mol, numbering: Dict[int, int], 
                                 highlight_atoms: List[int] = None,
                                 figsize=(10, 10)) -> None:
    """
    Draw molecule with ring atom numbering.
    
    Args:
        mol: RDKit molecule object
        numbering: Ring numbering dictionary
        highlight_atoms: Optional list of atom indices to highlight
        figsize: Figure size tuple
    """
    # Create atom labels
    atom_labels = {}
    for atom_idx, position in numbering.items():
        atom_labels[atom_idx] = str(position)
    
    # Draw with labels
    img = Draw.MolToImage(
        mol,
        size=(figsize[0]*100, figsize[1]*100),
        highlightAtoms=highlight_atoms or [],
        atomLabels=atom_labels
    )
    
    plt.figure(figsize=figsize)
    plt.imshow(img)
    plt.axis('off')
    plt.tight_layout()
    plt.show()


def plot_substitution_frequency(position_counts: Dict[int, int], 
                                save_path: str = None,
                                figsize=(12, 6)) -> None:
    """
    Plot bar chart of substitution frequency at each ring position.
    
    Args:
        position_counts: Dictionary mapping ring position to count
        save_path: Optional path to save the figure
        figsize: Figure size tuple
    """
    positions = sorted(position_counts.keys())
    counts = [position_counts[pos] for pos in positions]
    
    plt.figure(figsize=figsize)
    
    # Create bar plot
    bars = plt.bar(positions, counts, color='steelblue', edgecolor='black', alpha=0.7)
    
    # Color bars by frequency (gradient)
    max_count = max(counts) if counts else 1
    for bar, count in zip(bars, counts):
        normalized = count / max_count
        bar.set_color(plt.cm.YlOrRd(normalized))
    
    plt.xlabel('Ring Position', fontsize=14, fontweight='bold')
    plt.ylabel('Number of Substitutions', fontsize=14, fontweight='bold')
    plt.title('Side Chain Substitution Frequency by Ring Position', 
             fontsize=16, fontweight='bold')
    plt.xticks(positions, fontsize=12)
    plt.yticks(fontsize=12)
    plt.grid(axis='y', alpha=0.3, linestyle='--')
    
    # Add value labels on bars
    for i, (pos, count) in enumerate(zip(positions, counts)):
        plt.text(pos, count + max_count*0.01, str(count), 
                ha='center', va='bottom', fontsize=10, fontweight='bold')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Figure saved to {save_path}")
    
    plt.show()


def plot_substitution_heatmap(position_fragment_data: pd.DataFrame,
                              save_path: str = None,
                              figsize=(14, 8)) -> None:
    """
    Plot heatmap showing fragment diversity at each position.
    
    Args:
        position_fragment_data: DataFrame with position and fragment information
        save_path: Optional path to save the figure
        figsize: Figure size tuple
    """
    # Create pivot table
    pivot = position_fragment_data.pivot_table(
        index='fragment_smiles',
        columns='cleavage_position',
        values='count',
        fill_value=0
    )
    
    plt.figure(figsize=figsize)
    
    sns.heatmap(pivot, cmap='YlOrRd', annot=False, 
                cbar_kws={'label': 'Frequency'},
                linewidths=0.5, linecolor='gray')
    
    plt.xlabel('Ring Position', fontsize=14, fontweight='bold')
    plt.ylabel('Fragment Type', fontsize=14, fontweight='bold')
    plt.title('Fragment Diversity Across Ring Positions', 
             fontsize=16, fontweight='bold')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Figure saved to {save_path}")
    
    plt.show()


def plot_circular_substitution_pattern(position_counts: Dict[int, int],
                                       save_path: str = None,
                                       figsize=(10, 10)) -> None:
    """
    Plot circular/radar plot of substitution patterns around the ring.
    
    Args:
        position_counts: Dictionary mapping ring position to count
        save_path: Optional path to save the figure
        figsize: Figure size tuple
    """
    positions = list(range(1, 17))  # All 16 positions
    counts = [position_counts.get(pos, 0) for pos in positions]
    
    # Append first value to close the circle
    counts_circle = counts + [counts[0]]
    
    # Calculate angles
    angles = np.linspace(0, 2 * np.pi, 16, endpoint=False).tolist()
    angles_circle = angles + [angles[0]]
    
    fig, ax = plt.subplots(figsize=figsize, subplot_kw=dict(projection='polar'))
    
    # Plot
    ax.plot(angles_circle, counts_circle, 'o-', linewidth=2, color='steelblue', 
            markersize=8, markerfacecolor='orange', markeredgecolor='black', markeredgewidth=1.5)
    ax.fill(angles_circle, counts_circle, alpha=0.25, color='steelblue')
    
    # Set labels
    ax.set_xticks(angles)
    ax.set_xticklabels([str(i) for i in positions], fontsize=12)
    ax.set_ylim(0, max(counts) * 1.1 if counts else 1)
    
    ax.set_title('Circular Substitution Pattern\n(16-Membered Ring)', 
                fontsize=16, fontweight='bold', pad=20)
    ax.grid(True, linestyle='--', alpha=0.6)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Figure saved to {save_path}")
    
    plt.show()


def plot_fragment_statistics(fragments_df: pd.DataFrame,
                             save_path: str = None,
                             figsize=(15, 5)) -> None:
    """
    Plot statistics about fragment properties.
    
    Args:
        fragments_df: DataFrame with fragment information
        save_path: Optional path to save the figure
        figsize: Figure size tuple
    """
    fig, axes = plt.subplots(1, 3, figsize=figsize)
    
    # Plot 1: Fragment size distribution
    axes[0].hist(fragments_df['atom_count'], bins=20, color='skyblue', 
                edgecolor='black', alpha=0.7)
    axes[0].set_xlabel('Number of Atoms', fontsize=12, fontweight='bold')
    axes[0].set_ylabel('Frequency', fontsize=12, fontweight='bold')
    axes[0].set_title('Fragment Size Distribution', fontsize=14, fontweight='bold')
    axes[0].grid(axis='y', alpha=0.3)
    
    # Plot 2: Molecular weight distribution
    axes[1].hist(fragments_df['molecular_weight'], bins=20, color='lightcoral', 
                edgecolor='black', alpha=0.7)
    axes[1].set_xlabel('Molecular Weight (Da)', fontsize=12, fontweight='bold')
    axes[1].set_ylabel('Frequency', fontsize=12, fontweight='bold')
    axes[1].set_title('Fragment MW Distribution', fontsize=14, fontweight='bold')
    axes[1].grid(axis='y', alpha=0.3)
    
    # Plot 3: Fragments per molecule
    frags_per_mol = fragments_df.groupby('parent_id').size()
    axes[2].hist(frags_per_mol, bins=range(int(frags_per_mol.min()), 
                                           int(frags_per_mol.max())+2), 
                color='lightgreen', edgecolor='black', alpha=0.7)
    axes[2].set_xlabel('Fragments per Molecule', fontsize=12, fontweight='bold')
    axes[2].set_ylabel('Number of Molecules', fontsize=12, fontweight='bold')
    axes[2].set_title('Fragments per Macrolactone', fontsize=14, fontweight='bold')
    axes[2].grid(axis='y', alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Figure saved to {save_path}")
    
    plt.show()

