import squidpy as sq
import scanpy as sc
import pandas as pd
import numpy as np

# Requires spatial neighbors
sq.gr.spatial_neighbors(adata, coord_type='generic', n_neighs=6)

# Compute Moran's I for all genes (can be slow)
sq.gr.spatial_autocorr(adata, mode='moran')

# Or for specific genes
sq.gr.spatial_autocorr(adata, mode='moran', genes=['GENE1', 'GENE2', 'GENE3'])

# Results stored in adata.uns['moranI']
moran_results = adata.uns['moranI']
print(moran_results.head(20))

# Moran's I ranges from -1 to 1
# I > 0: positive spatial autocorrelation (similar values cluster)
# I = 0: random spatial distribution
# I < 0: negative spatial autocorrelation (dissimilar values cluster)

# Get significantly spatially variable genes
svg = moran_results[moran_results['pval_norm'] < 0.05].sort_values('I', ascending=False)
print(f'Found {len(svg)} spatially variable genes (p < 0.05)')
print('\nTop 10 spatially variable genes:')
print(svg.head(10)[['I', 'pval_norm']])

# Alternative spatial autocorrelation measure
sq.gr.spatial_autocorr(adata, mode='geary')

# Results in adata.uns['gearyC']
geary_results = adata.uns['gearyC']
# C < 1: positive spatial autocorrelation
# C = 1: random
# C > 1: negative spatial autocorrelation

# Analyze co-localization of cell types/clusters
# First, ensure you have cluster labels
sc.pp.neighbors(adata)
sc.tl.leiden(adata)

# Compute co-occurrence
sq.gr.co_occurrence(adata, cluster_key='leiden')

# Results in adata.uns['leiden_co_occurrence']
# Visualize co-occurrence
sq.pl.co_occurrence(adata, cluster_key='leiden')

co_occ = adata.uns['leiden_co_occurrence']
occ_matrix = co_occ['occ']  # Occurrence matrix
interval = co_occ['interval']  # Distance intervals

# occ_matrix[i, j, k] = occurrence of cluster j around cluster i at distance interval k
print(f'Occurrence matrix shape: {occ_matrix.shape}')
print(f'Distance intervals: {interval}')

# Test if clusters are enriched in each other's neighborhoods
sq.gr.nhood_enrichment(adata, cluster_key='leiden')

# Results in adata.uns['leiden_nhood_enrichment']
# zscore > 0: clusters co-localize more than expected
# zscore < 0: clusters avoid each other

# Visualize
sq.pl.nhood_enrichment(adata, cluster_key='leiden')

enrichment = adata.uns['leiden_nhood_enrichment']
zscore = enrichment['zscore']
clusters = adata.obs['leiden'].cat.categories

# Convert to DataFrame
zscore_df = pd.DataFrame(zscore, index=clusters, columns=clusters)
print('Neighborhood enrichment z-scores:')
print(zscore_df)

# Ripley's K/L function for point pattern analysis (single-cell resolution data)
sq.gr.ripley(adata, cluster_key='leiden', mode='L')

# Results in adata.uns['leiden_ripley']
sq.pl.ripley(adata, cluster_key='leiden')

# Compute centrality of each cell type
sq.gr.centrality_scores(adata, cluster_key='leiden')

# Results in adata.uns['leiden_centrality_scores']
centrality = adata.uns['leiden_centrality_scores']
print(centrality)

# Build interaction matrix between clusters
sq.gr.interaction_matrix(adata, cluster_key='leiden')

# Results in adata.uns['leiden_interactions']
interactions = adata.uns['leiden_interactions']
print(interactions)

from scipy.stats import pearsonr

def spatial_correlation(adata, gene1, gene2):
    '''Compute spatial correlation between two genes'''
    expr1 = adata[:, gene1].X.toarray().flatten()
    expr2 = adata[:, gene2].X.toarray().flatten()
    r, p = pearsonr(expr1, expr2)
    return r, p

r, p = spatial_correlation(adata, 'GENE1', 'GENE2')
print(f'Spatial correlation: r={r:.3f}, p={p:.2e}')

from esda.moran import Moran_Local
from libpysal.weights import KNN

# Build weights matrix
coords = adata.obsm['spatial']
w = KNN.from_array(coords, k=6)
w.transform = 'r'

# Compute local Moran's I for a gene
gene_expr = adata[:, 'GENE1'].X.toarray().flatten()
lisa = Moran_Local(gene_expr, w)

# Add to adata
adata.obs['GENE1_lisa'] = lisa.Is
adata.obs['GENE1_lisa_q'] = lisa.q  # Quadrant (HH, HL, LH, LL)

# Compute Moran's I for top variable genes only
hvg = adata.var_names[adata.var['highly_variable']][:500]
sq.gr.spatial_autocorr(adata, mode='moran', genes=hvg)

results = adata.uns['moranI']
significant = results[results['pval_norm'] < 0.01]
print(f'{len(significant)} genes with significant spatial autocorrelation')