import gradio as gr
import scanpy as sc
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
import warnings
import io
import base64
from datetime import datetime

warnings.filterwarnings('ignore')

# تنظیمات اولیه
plt.style.use('default')
sc.settings.verbosity = 1
sc.settings.set_figure_params(dpi=100, facecolor='white')

class SingleCellAnalyzer:
    def __init__(self):
        self.adata = None
        self.results = {}
        
    def load_dataset(self, dataset_name):
        """بارگذاری dataset"""
        try:
            if dataset_name == "PBMC 3k":
                self.adata = sc.datasets.pbmc3k()
            elif dataset_name == "PBMC 68k":
                self.adata = sc.datasets.pbmc68k_reduced()
            elif dataset_name == "Paul15":
                self.adata = sc.datasets.paul15()
            elif dataset_name == "Mouse Brain":
                # بارگذاری از URL
                url = "https://ndownloader.figshare.com/files/24539928"
                self.adata = sc.read_h5ad(url, backup_url=url)
            else:
                return False, "Dataset not found"
            
            return True, f"Dataset loaded: {self.adata.shape[0]} cells, {self.adata.shape[1]} genes"
        except Exception as e:
            return False, f"Error loading dataset: {str(e)}"
    
    def preprocess_data(self, min_genes=200, min_cells=3, target_sum=1e4):
        """پیش‌پردازش داده‌ها"""
        if self.adata is None:
            return False, "No dataset loaded"
        
        try:
            # فیلتر سلول‌ها و ژن‌ها
            sc.pp.filter_cells(self.adata, min_genes=min_genes)
            sc.pp.filter_genes(self.adata, min_cells=min_cells)
            
            # محاسبه QC metrics
            self.adata.var['mt'] = self.adata.var_names.str.startswith('MT-')
            sc.pp.calculate_qc_metrics(
                self.adata, 
                qc_vars=['mt'], 
                percent_top=None, 
                log1p=False, 
                inplace=True
            )
            
            # فیلتر کردن outlierها
            self.adata = self.adata[self.adata.obs.n_genes_by_counts < 5000, :]
            self.adata = self.adata[self.adata.obs.pct_counts_mt < 10, :]
            
            # نرمال‌سازی
            sc.pp.normalize_total(self.adata, target_sum=target_sum)
            sc.pp.log1p(self.adata)
            
            # شناسایی ژن‌های highly variable
            sc.pp.highly_variable_genes(
                self.adata, 
                min_mean=0.0125, 
                max_mean=3, 
                min_disp=0.5
            )
            
            # تنظیم مقیاس
            self.adata = self.adata[:, self.adata.var.highly_variable]
            sc.pp.regress_out(self.adata, ['total_counts', 'pct_counts_mt'])
            sc.pp.scale(self.adata, max_value=10)
            
            return True, "Preprocessing completed successfully"
        except Exception as e:
            return False, f"Error in preprocessing: {str(e)}"
    
    def run_analysis(self, n_neighbors=15, n_pcs=40, resolution=0.5):
        """اجرای تحلیل کامل"""
        if self.adata is None:
            return False, "No data available for analysis"
        
        try:
            # PCA
            sc.tl.pca(self.adata, svd_solver='arpack', n_comps=min(50, self.adata.shape[1]-1))
            
            # همسایگی
            sc.pp.neighbors(self.adata, n_neighbors=n_neighbors, n_pcs=min(n_pcs, self.adata.shape[1]-1))
            
            # UMAP
            sc.tl.umap(self.adata)
            
            # t-SNE
            sc.tl.tsne(self.adata, n_pcs=min(20, self.adata.shape[1]-1))
            
            # خوشه‌بندی
            sc.tl.leiden(self.adata, resolution=resolution, key_added='clusters')
            
            # مارکرهای هر خوشه
            sc.tl.rank_genes_groups(
                self.adata, 
                'clusters', 
                method='t-test',
                key_added='rank_genes'
            )
            
            # ذخیره نتایج
            self.results = {
                'n_cells': self.adata.shape[0],
                'n_genes': self.adata.shape[1],
                'n_clusters': len(self.adata.obs['clusters'].unique()),
                'clusters': self.adata.obs['clusters'].value_counts().to_dict(),
                'top_markers': self.get_top_markers(5)
            }
            
            return True, f"Analysis complete: Found {self.results['n_clusters']} clusters"
        except Exception as e:
            return False, f"Error in analysis: {str(e)}"
    
    def get_top_markers(self, n_genes=5):
        """استخراج مارکرهای برتر"""
        if 'rank_genes' not in self.adata.uns:
            return {}
        
        markers = {}
        result = self.adata.uns['rank_genes']
        
        for cluster in result['names'].dtype.names:
            cluster_markers = []
            for i in range(min(n_genes, len(result['names'][cluster]))):
                gene = result['names'][cluster][i]
                score = result['scores'][cluster][i]
                pval = result['pvals'][cluster][i] if 'pvals' in result else None
                
                cluster_markers.append({
                    'gene': gene,
                    'score': float(score),
                    'p_value': float(pval) if pval else None
                })
            markers[cluster] = cluster_markers
        
        return markers
    
    def create_visualizations(self):
        """ایجاد تجسم‌ها"""
        if self.adata is None:
            return None, None, None, None
        
        try:
            # 1. UMAP plot با plotly
            umap_df = pd.DataFrame({
                'UMAP1': self.adata.obsm['X_umap'][:, 0],
                'UMAP2': self.adata.obsm['X_umap'][:, 1],
                'Cluster': self.adata.obs['clusters'],
                'n_genes': self.adata.obs['n_genes_by_counts']
            })
            
            fig_umap = px.scatter(
                umap_df, 
                x='UMAP1', 
                y='UMAP2', 
                color='Cluster',
                size='n_genes',
                hover_data=['n_genes'],
                title='UMAP Visualization',
                color_discrete_sequence=px.colors.qualitative.Set3
            )
            fig_umap.update_layout(height=500)
            
            # 2. نمودار توزیع خوشه‌ها
            cluster_counts = self.adata.obs['clusters'].value_counts().sort_index()
            fig_clusters = go.Figure(data=[
                go.Bar(
                    x=list(cluster_counts.index),
                    y=cluster_counts.values,
                    marker_color=px.colors.qualitative.Set3[:len(cluster_counts)],
                    text=cluster_counts.values,
                    textposition='auto'
                )
            ])
            fig_clusters.update_layout(
                title='Cell Distribution per Cluster',
                xaxis_title='Cluster',
                yaxis_title='Number of Cells',
                height=400
            )
            
            # 3. نمودار کیفیت داده‌ها
            fig_qc = go.Figure()
            
            fig_qc.add_trace(go.Violin(
                y=self.adata.obs['n_genes_by_counts'],
                box_visible=True,
                line_color='blue',
                name='Genes per Cell'
            ))
            
            fig_qc.add_trace(go.Violin(
                y=self.adata.obs['total_counts'],
                box_visible=True,
                line_color='red',
                name='Total Counts'
            ))
            
            fig_qc.update_layout(
                title='Data Quality Metrics',
                yaxis_title='Value',
                height=400
            )
            
            # 4. Heatmap مارکرها
            # استخراج 3 مارکر برتر از هر خوشه
            top_marker_genes = []
            markers_dict = self.get_top_markers(3)
            for cluster in markers_dict:
                for marker in markers_dict[cluster]:
                    top_marker_genes.append(marker['gene'])
            
            # محدود کردن به 30 ژن
            top_marker_genes = list(set(top_marker_genes))[:30]
            
            # ایجاد heatmap
            if len(top_marker_genes) > 0:
                # میانگین اکسپرشن برای هر خوشه
                expr_matrix = []
                cluster_order = sorted(self.adata.obs['clusters'].unique())
                
                for gene in top_marker_genes:
                    if gene in self.adata.var_names:
                        gene_expr = []
                        for cluster in cluster_order:
                            mask = self.adata.obs['clusters'] == cluster
                            mean_expr = self.adata[mask, gene].X.mean()
                            gene_expr.append(mean_expr)
                        expr_matrix.append(gene_expr)
                
                if expr_matrix:
                    fig_heatmap = go.Figure(data=go.Heatmap(
                        z=expr_matrix,
                        x=cluster_order,
                        y=top_marker_genes[:len(expr_matrix)],
                        colorscale='Viridis',
                        colorbar=dict(title='Mean Expression')
                    ))
                    fig_heatmap.update_layout(
                        title='Top Marker Genes Expression',
                        xaxis_title='Cluster',
                        yaxis_title='Gene',
                        height=500
                    )
                else:
                    fig_heatmap = go.Figure()
                    fig_heatmap.update_layout(
                        title='No marker genes found',
                        height=300
                    )
            else:
                fig_heatmap = go.Figure()
                fig_heatmap.update_layout(
                    title='No marker genes found',
                    height=300
                )
            
            return fig_umap, fig_clusters, fig_qc, fig_heatmap
            
        except Exception as e:
            print(f"Visualization error: {e}")
            return None, None, None, None

# ایجاد instance از آنالایزر
analyzer = SingleCellAnalyzer()

# تابع اصلی برای Gradio
def process_pipeline(dataset_name, min_genes, min_cells, n_neighbors, resolution):
    """پایپ‌لاین پردازش"""
    messages = []
    
    # بارگذاری داده‌ها
    success, msg = analyzer.load_dataset(dataset_name)
    messages.append(f"📥 **Data Loading**: {msg}")
    if not success:
        return None, None, None, None, "\n".join(messages), {}
    
    # پیش‌پردازش
    success, msg = analyzer.preprocess_data(
        min_genes=int(min_genes), 
        min_cells=int(min_cells)
    )
    messages.append(f"🔧 **Preprocessing**: {msg}")
    if not success:
        return None, None, None, None, "\n".join(messages), {}
    
    # تحلیل
    success, msg = analyzer.run_analysis(
        n_neighbors=int(n_neighbors),
        resolution=float(resolution)
    )
    messages.append(f"🔬 **Analysis**: {msg}")
    
    # تولید تجسم‌ها
    fig_umap, fig_clusters, fig_qc, fig_heatmap = analyzer.create_visualizations()
    
    # آمار نتایج
    stats = analyzer.results if analyzer.results else {}
    
    return fig_umap, fig_clusters, fig_qc, fig_heatmap, "\n".join(messages), stats

# رابط Gradio
with gr.Blocks(
    title="Single-Cell RNA-seq Analyzer",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="teal"),
    css="""
    .gradio-container {max-width: 1200px !important;}
    .message-box {background-color: #f0f8ff; padding: 15px; border-radius: 10px;}
    """
) as demo:
    
    gr.Markdown("""
    # 🔬 Single-Cell RNA-seq Analysis Platform
    ### Interactive tool for clustering and visualizing single-cell data
    
    *This tool performs automated analysis of single-cell RNA sequencing data including 
    preprocessing, dimensionality reduction, clustering, and marker gene identification.*
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Analysis Parameters")
            
            dataset = gr.Dropdown(
                choices=["PBMC 3k", "PBMC 68k", "Paul15", "Mouse Brain"],
                value="PBMC 3k",
                label="Select Dataset",
                info="Choose a pre-loaded single-cell dataset"
            )
            
            with gr.Row():
                min_genes = gr.Number(
                    value=200,
                    label="Minimum Genes per Cell",
                    minimum=100,
                    maximum=1000
                )
                
                min_cells = gr.Number(
                    value=3,
                    label="Minimum Cells per Gene",
                    minimum=1,
                    maximum=10
                )
            
            with gr.Row():
                n_neighbors = gr.Slider(
                    minimum=5,
                    maximum=50,
                    value=15,
                    step=1,
                    label="Neighbors for Graph"
                )
                
                resolution = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.5,
                    step=0.1,
                    label="Clustering Resolution"
                )
            
            analyze_btn = gr.Button(
                "🚀 Run Analysis",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("---")
            gr.Markdown("### 📊 Quick Statistics")
            stats_output = gr.JSON(label="", value={})
        
        with gr.Column(scale=2):
            gr.Markdown("### 📈 Visualizations")
            
            with gr.Tab("UMAP Plot"):
                umap_plot = gr.Plot(label="")
            
            with gr.Tab("Cluster Distribution"):
                cluster_plot = gr.Plot(label="")
            
            with gr.Tab("Quality Control"):
                qc_plot = gr.Plot(label="")
            
            with gr.Tab("Marker Genes"):
                heatmap_plot = gr.Plot(label="")
    
    # بخش لاگ و پیام‌ها
    with gr.Row():
        message_output = gr.Markdown(
            "### 📋 Analysis Log\n*Ready to analyze...*",
            elem_classes="message-box"
        )
    
    # مثال‌ها
    gr.Markdown("### 🚀 Quick Start Examples")
    examples = gr.Examples(
        examples=[
            ["PBMC 3k", 200, 3, 15, 0.5],
            ["PBMC 68k", 300, 5, 20, 0.4],
            ["Paul15", 250, 4, 12, 0.6]
        ],
        inputs=[dataset, min_genes, min_cells, n_neighbors, resolution],
        outputs=[umap_plot, cluster_plot, qc_plot, heatmap_plot, message_output, stats_output],
        fn=process_pipeline,
        cache_examples=False
    )
    
    # اتصال دکمه
    analyze_btn.click(
        fn=process_pipeline,
        inputs=[dataset, min_genes, min_cells, n_neighbors, resolution],
        outputs=[umap_plot, cluster_plot, qc_plot, heatmap_plot, message_output, stats_output]
    )

# اجرای برنامه
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )