import numpy as np


def load_expression_matrix(path: str) -> tuple[np.ndarray, list[str], list[str]]:
    """
    반환: (matrix, gene_names, sample_names)
    matrix shape: (n_genes, n_samples)
    """
    with open(path) as f:
        header = f.readline().strip().split("\t")
        sample_names = header[1:]
        
        gene_names = []
        rows = []
        for line in f:
            parts = line.strip().split("\t")
            gene_names.append(parts[0])
            rows.append([float(x) for x in parts[1:]])
    
    return np.array(rows), gene_names, sample_names


expr, genes, samples = load_expression_matrix("expression.tsv")
print(f"유전자 {len(genes)}, 샘플 {len(samples)}")
print(f"행렬 shape: {expr.shape}")

def compute_correlation_matrix(expr: np.ndarray) -> np.ndarray:
    """
    expr shape: (n_genes, n_samples)
    반환: (n_genes, n_genes) 상관관계 행렬
    """
    return np.corrcoef(expr)


corr = compute_correlation_matrix(expr)
print(f"corr shape: {corr.shape}")  # (5000, 5000)
print(f"대각선: {corr[0, 0]}")  # 1.0 (자기 자신)

def correlation_to_distance(corr: np.ndarray) -> np.ndarray:
    """
    거리 = 1 - correlation.
    양의 상관 강함 → 거리 작음.
    음의 상관 (반대로 움직임) → 거리 큼.
    """
    return 1.0 - corr


dist = correlation_to_distance(corr)

def hierarchical_clustering_single(dist: np.ndarray) -> list[tuple[int, int, float]]:
    """
    반환: 병합 이벤트 리스트 [(cluster_a, cluster_b, merge_distance), ...]
    각 병합 후 새 클러스터 id는 원래 개수 + 이벤트 인덱스로 부여.
    """
    n = dist.shape[0]
    active_clusters = {i: [i] for i in range(n)}
    cluster_distances = {(i, j): dist[i, j] for i in range(n) for j in range(i + 1, n)}
    
    events: list[tuple[int, int, float]] = []
    next_id = n
    
    while len(active_clusters) > 1:
        best_pair = min(cluster_distances, key=cluster_distances.get)
        i, j = best_pair
        merge_dist = cluster_distances[best_pair]
        events.append((i, j, merge_dist))
        
        new_cluster = active_clusters[i] + active_clusters[j]
        del active_clusters[i]
        del active_clusters[j]
        active_clusters[next_id] = new_cluster
        
        # 새 클러스터와 나머지 클러스터의 거리 계산 (single linkage = min)
        new_distances = {}
        for existing_id in active_clusters:
            if existing_id == next_id:
                continue
            existing = active_clusters[existing_id]
            min_d = min(
                dist[a, b] for a in new_cluster for b in existing
            )
            new_distances[(min(existing_id, next_id), max(existing_id, next_id))] = min_d
        
        cluster_distances = {
            k: v for k, v in cluster_distances.items()
            if i not in k and j not in k
        }
        cluster_distances.update(new_distances)
        
        next_id += 1
    
    return events

def cut_dendrogram(events: list[tuple[int, int, float]], n_leaves: int, num_clusters: int):
    parent = list(range(n_leaves + len(events)))
    
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    
    def union(x, y, new_id):
        rx, ry = find(x), find(y)
        parent[rx] = new_id
        parent[ry] = new_id
    
    # 마지막 (num_clusters - 1) 개의 병합만 무시하고 그 앞까지 union
    events_to_apply = events[:-num_clusters + 1] if num_clusters > 1 else events
    
    next_id = n_leaves
    for a, b, _ in events_to_apply:
        union(a, b, next_id)
        next_id += 1
    
    clusters = {}
    for leaf in range(n_leaves):
        root = find(leaf)
        clusters.setdefault(root, []).append(leaf)
    
    return list(clusters.values())


cluster_lists = cut_dendrogram(events, n_leaves=len(genes), num_clusters=10)
for i, members in enumerate(cluster_lists):
    print(f"클러스터 {i}: {len(members)} 유전자")

def cluster_order(cluster_lists: list[list[int]]) -> list[int]:
    """각 클러스터 안에서 원래 인덱스 순서를 유지하며, 클러스터별로 정렬."""
    result = []
    for members in cluster_lists:
        result.extend(sorted(members))
    return result


new_order = cluster_order(cluster_lists)
reordered_expr = expr[new_order]

import matplotlib.pyplot as plt


def plot_heatmap(matrix: np.ndarray, ax=None) -> None:
    if ax is None:
        _, ax = plt.subplots(figsize=(6, 8))
    im = ax.imshow(matrix, cmap="RdBu_r", aspect="auto")
    plt.colorbar(im, ax=ax)


fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))
plot_heatmap(expr, ax1)
ax1.set_title("Original")
plot_heatmap(reordered_expr, ax2)
ax2.set_title("Clustered")
plt.tight_layout()
plt.show()

def hierarchical_clustering_complete(dist: np.ndarray) -> list[tuple[int, int, float]]:
    # single linkage 코드와 대부분 동일, min → max 로 교체
    # ...
    # 클러스터간 거리 계산 부분:
    # TODO: min(dist[a, b] for ...) 를 max(dist[a, b] for ...) 로
    pass

def find_hub_genes(
    corr: np.ndarray,
    clusters: list[list[int]],
    gene_names: list[str]
) -> list[tuple[str, float]]:
    """각 클러스터의 대표 유전자와 그 평균 상관관계."""
    hubs = []
    for members in clusters:
        # TODO: 각 멤버에 대해 다른 멤버들과의 평균 상관관계 계산
        # 가장 높은 값을 가진 유전자를 hub로 선정
        pass
    return hubs

best_score = -1
best_gene = None
for m in members:
    score = np.mean([corr[m, other] for other in members if other != m])
    if score > best_score:
        best_score = score
        best_gene = gene_names[m]

def enrich_clusters(
    clusters: list[list[str]],
    gene_to_pathway_csv: str
) -> dict:
    """
    각 클러스터에서 어떤 pathway가 가장 많이 나타나는지 상위 3개.
    """
    # TODO: CSV 로드 (gene, pathway 컬럼)
    # 각 클러스터에서 pathway 등장 빈도 계산
    # 각 클러스터의 상위 3 pathway 반환
    pass