from heartmap import Config
from heartmap.pipelines import ComprehensivePipeline
import scanpy as sc
import pandas as pd
from pathlib import Path

# ========================================
# Setup
# ========================================
print("=== HeartMAP Comprehensive Analysis ===")

# Configure
config = Config.default()
config.data.min_genes = 200
config.data.min_cells = 3
config.data.max_cells_subset = 50000
config.data.n_top_genes = 2000
config.analysis.resolution = 0.5
config.analysis.n_neighbors = 10
config.analysis.n_pcs = 40
config.analysis.use_liana = True

config.update_paths('./analysis')
config.create_directories()

# ========================================
# Run Pipeline
# ========================================
pipeline = ComprehensivePipeline(config)
results = pipeline.run(
    data_path='data/raw/heart_data.h5ad',
    output_dir='results/comprehensive'
)

# ========================================
# Analyze Results
# ========================================
adata = results['adata']
analysis_results = results['results']

print(f"\n{'='*50}")
print("ANALYSIS SUMMARY")
print('='*50)

print(f"\nDataset: {adata.n_obs:,} cells × {adata.n_vars:,} genes")

# Cell type annotation
if 'leiden' in adata.obs.columns:
    n_clusters = len(adata.obs['leiden'].unique())
    print(f"\nCell Type Clusters: {n_clusters}")
    
    cluster_dist = adata.obs['leiden'].value_counts().head(5)
    print("\nTop 5 Clusters:")
    for cluster, count in cluster_dist.items():
        pct = 100 * count / adata.n_obs
        print(f"  Cluster {cluster}: {count:,} cells ({pct:.1f}%)")

# Chamber distribution
if 'chamber' in adata.obs.columns:
    print("\nChamber Distribution:")
    for chamber, count in adata.obs['chamber'].value_counts().items():
        pct = 100 * count / adata.n_obs
        print(f"  {chamber}: {count:,} cells ({pct:.1f}%)")

# Communication hubs
if 'hub_score' in adata.obs.columns:
    hub_mean = adata.obs['hub_score'].mean()
    hub_std = adata.obs['hub_score'].std()
    print(f"\nHub Scores: {hub_mean:.4f} ± {hub_std:.4f}")
    
    # Top hub clusters
    hub_by_cluster = adata.obs.groupby('leiden')['hub_score'].mean()
    top_hubs = hub_by_cluster.nlargest(3)
    print("\nTop Communication Hub Clusters:")
    for cluster, score in top_hubs.items():
        print(f"  Cluster {cluster}: {score:.4f}")

# ========================================
# Generate Custom Visualizations
# ========================================
print(f"\n{'='*50}")
print("GENERATING VISUALIZATIONS")
print('='*50)

import matplotlib.pyplot as plt

# Multi-panel UMAP
if 'chamber' in adata.obs.columns:
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    
    # Panel 1: Clusters
    sc.pl.umap(adata, color='leiden', ax=axes[0], 
               show=False, frameon=False, title='Cell Types')
    
    # Panel 2: Chambers
    sc.pl.umap(adata, color='chamber', ax=axes[1],
               show=False, frameon=False, title='Chambers')
    
    # Panel 3: Hub scores
    if 'hub_score' in adata.obs.columns:
        sc.pl.umap(adata, color='hub_score', ax=axes[2],
                   show=False, frameon=False, title='Hub Scores',
                   cmap='viridis')
    
    plt.tight_layout()
    plt.savefig('results/comprehensive/overview_panel.png',
                dpi=300, bbox_inches='tight')
    plt.close()
    print("  Saved: overview_panel.png")

# ========================================
# Export Results
# ========================================
print(f"\n{'='*50}")
print("EXPORTING RESULTS")
print('='*50)

# Cell metadata
metadata = adata.obs[['leiden', 'n_genes', 'total_counts']].copy()
if 'chamber' in adata.obs.columns:
    metadata['chamber'] = adata.obs['chamber']
if 'hub_score' in adata.obs.columns:
    metadata['hub_score'] = adata.obs['hub_score']

metadata.to_csv('results/comprehensive/cell_metadata.csv')
print("  Saved: cell_metadata.csv")

# Summary statistics
summary = {
    'total_cells': adata.n_obs,
    'total_genes': adata.n_vars,
    'n_clusters': len(adata.obs['leiden'].unique()),
    'mean_genes_per_cell': adata.obs['n_genes'].mean(),
    'mean_counts_per_cell': adata.obs['total_counts'].mean()
}

if 'chamber' in adata.obs.columns:
    summary['n_chambers'] = len(adata.obs['chamber'].unique())

summary_df = pd.DataFrame([summary])
summary_df.to_csv('results/comprehensive/summary_stats.csv', index=False)
print("  Saved: summary_stats.csv")

print(f"\n{'='*50}")
print("ANALYSIS COMPLETE!")
print('='*50)
print("\nOutput files:")
for f in Path('results/comprehensive').rglob('*'):
    if f.is_file():
        print(f"  {f.relative_to('results/comprehensive')}")

results/comprehensive/
├── heartmap_complete.h5ad           # Complete processed data
├── heartmap_model.pkl               # Trained model (if saved)
├── analysis_report.md               # Automated report
├── cell_metadata.csv                # All cell annotations
├── summary_stats.csv                # Dataset statistics
├── figures/
│   ├── comprehensive_dashboard.png  # Integrated visualization
│   ├── umap_clusters.png           # Cell type UMAP
│   ├── qc_metrics.png              # Quality control plots
│   ├── communication_heatmap.png   # Communication matrix
│   ├── hub_scores.png              # Hub score UMAP
│   ├── chamber_composition.png     # Chamber distribution
│   └── chamber_correlations.png    # Cross-chamber similarity
└── tables/
    ├── marker_genes.csv            # Cell type markers
    ├── communication_scores.csv    # L-R interactions
    ├── hub_scores_by_type.csv     # Hub scores per type
    └── chamber_markers.csv         # Chamber-specific genes