import argparse
import sys
from typing import TextIO

import cellxgene_census
import pandas as pd
from tiledbsoma import DataFrame

from .build_soma.globals import CENSUS_DATA_NAME, CENSUS_INFO_NAME, CENSUS_SPATIAL_SEQUENCING_NAME

# Print all of the Pandas DataFrames, except the dimensions
pd.options.display.max_columns = None
pd.options.display.max_rows = None
pd.options.display.width = 1024
pd.options.display.show_dimensions = False


def display_summary(
    *,
    census_version: str | None = "latest",
    uri: str | None = None,
    file: TextIO | None = None,
) -> int:
    census = cellxgene_census.open_soma(census_version=census_version, uri=uri)

    COLS_TO_QUERY = [
        ("soma_joinid", "cells"),
        ("dataset_id", "datasets"),
        ("cell_type_ontology_term_id", "cell types"),
        ("tissue_ontology_term_id", "tissues"),
        ("assay_ontology_term_id", "assays"),
    ]

    def _read_collection_obs(collection_name: str) -> dict[tuple[str, str], pd.DataFrame]:
        results: dict[tuple[str, str], pd.DataFrame] = {}
        for name, experiment in census.get(collection_name).items():
            results[(name, collection_name)] = (
                experiment.obs.read(column_names=[c[0] for c in COLS_TO_QUERY]).concat().to_pandas()
            )
        return results

    # Read both collections and merge into a dict keyed by (experiment_name, collection_name)
    experiments_obs = {
        **_read_collection_obs(CENSUS_DATA_NAME),
        **_read_collection_obs(CENSUS_SPATIAL_SEQUENCING_NAME),
    }

    # Use Pandas to summarize and display
    # Build stats including collection column; compute unique counts per organism per collection
    stats = []
    for (experiment_name, collection_name), df in experiments_obs.items():
        for col in COLS_TO_QUERY:
            stats.append((experiment_name, collection_name, col[1], df[col[0]].nunique()))
    print(
        census["census_info"]["summary"].read().concat().to_pandas()[["label", "value"]].to_string(index=False),
        file=file,
    )
    stats_df = pd.DataFrame(stats, columns=["organism", "collection", "attribute", "unique count"])
    display_stats_df = stats_df.pivot_table(
        index=["organism", "collection"], columns=["attribute"], values=["unique count"], aggfunc="first"
    )
    print(display_stats_df, file=file)
    print(file=file)

    return 0


def display_diff(
    census_version: str | None = "latest",
    uri: str | None = None,
    previous_census_version: str | None = None,
    previous_uri: str | None = None,
    file: TextIO | None = None,
) -> int:
    census = cellxgene_census.open_soma(census_version=census_version, uri=uri)
    previous_census = cellxgene_census.open_soma(census_version=previous_census_version, uri=previous_uri)

    # Total cell count deltas by experiment (mouse, human)

    for organism in census[CENSUS_DATA_NAME]:
        curr_count = census[CENSUS_DATA_NAME][organism].obs.count
        if organism in previous_census[CENSUS_DATA_NAME]:
            prev_count = previous_census[CENSUS_DATA_NAME][organism].obs.count
            print(
                f"Previous {organism} cell count: {prev_count}, current {organism} cell count: {curr_count}, delta {curr_count - prev_count}",
                file=file,
            )
        else:
            print(
                f"New organism {organism} added, cell count: {curr_count}.",
                file=file,
            )
        print(file=file)

    prev_datasets = previous_census[CENSUS_INFO_NAME]["datasets"].read().concat().to_pandas()
    curr_datasets = census[CENSUS_INFO_NAME]["datasets"].read().concat().to_pandas()

    # Build dataset_id -> organism mapping for each census so we can report organism for added/removed datasets
    def _build_dataset_organism_map(census_obj: DataFrame) -> dict[str, str]:
        mapping: dict[str, set[str]] = {}
        for org, exp in census_obj[CENSUS_DATA_NAME].items():
            df = exp.obs.read(column_names=["dataset_id"]).concat().to_pandas()
            for dsid in df["dataset_id"].unique():
                mapping.setdefault(dsid, set()).add(org)
        # Convert sets to a stable, comma-separated string (handles rare multi-organism datasets)
        return {k: ", ".join(sorted(v)) for k, v in mapping.items()}

    curr_ds_to_org = _build_dataset_organism_map(census)
    prev_ds_to_org = _build_dataset_organism_map(previous_census)

    # Datasets removed, added
    curr_datasets_ids = set(curr_datasets["dataset_id"])
    prev_dataset_ids = set(prev_datasets["dataset_id"])

    added_datasets = curr_datasets_ids - prev_dataset_ids
    removed_datasets = prev_dataset_ids - curr_datasets_ids
    if added_datasets:
        print(f"Datasets that were added ({len(added_datasets)})", file=file)
        added_datasets_df = curr_datasets[curr_datasets["dataset_id"].isin(added_datasets)].copy()
        added_datasets_df["organism"] = added_datasets_df["dataset_id"].map(curr_ds_to_org).fillna("unknown")
        print(added_datasets_df[["dataset_id", "organism", "dataset_title", "collection_name"]], file=file)
    else:
        print("No datasets were added", file=file)
    print(file=file)

    if removed_datasets:
        print(f"Datasets that were removed ({len(removed_datasets)})", file=file)
        removed_datasets_df = prev_datasets[prev_datasets["dataset_id"].isin(removed_datasets)].copy()
        removed_datasets_df["organism"] = removed_datasets_df["dataset_id"].map(prev_ds_to_org).fillna("unknown")
        print(removed_datasets_df[["dataset_id", "organism", "dataset_title", "collection_name"]], file=file)
    else:
        print("No datasets were removed", file=file)
    print(file=file)

    # Datasets in both versions but that have differing cell counts
    joined = prev_datasets.join(
        curr_datasets.set_index("dataset_id"), on="dataset_id", lsuffix="_prev", rsuffix="_curr"
    )
    datasets_with_different_cell_counts = joined.loc[
        joined["dataset_total_cell_count_prev"] != joined["dataset_total_cell_count_curr"]
    ][["dataset_id", "dataset_total_cell_count_prev", "dataset_total_cell_count_curr"]]

    if not datasets_with_different_cell_counts.empty:
        org_series = datasets_with_different_cell_counts["dataset_id"].map(curr_ds_to_org)
        datasets_with_different_cell_counts.insert(1, "organism", org_series.fillna("???"))

        print("Datasets that have a different cell count", file=file)
        print(
            datasets_with_different_cell_counts[
                [
                    "dataset_id",
                    "organism",
                    "dataset_total_cell_count_prev",
                    "dataset_total_cell_count_curr",
                ]
            ],
            file=file,
        )
        print(file=file)

    # Deltas between summary_cell_counts dataframes
    y = census["census_info"]["summary_cell_counts"].read().concat().to_pandas()
    y = y.set_index(["organism", "category", "ontology_term_id"])

    z = previous_census["census_info"]["summary_cell_counts"].read().concat().to_pandas()
    z = z.set_index(["organism", "category", "ontology_term_id"])

    w = y.join(z, lsuffix="_prev", rsuffix="_curr")
    delta = w.loc[w["total_cell_count_prev"] != w["total_cell_count_curr"]][
        ["total_cell_count_prev", "total_cell_count_curr"]
    ].reset_index()
    if not delta.empty:
        print("Summary delta - total cell counts", file=file)
        print(delta, file=file)
        print(file=file)

    delta = w.loc[w["unique_cell_count_prev"] != w["unique_cell_count_curr"]][
        ["unique_cell_count_prev", "unique_cell_count_curr"]
    ].reset_index()
    if not delta.empty:
        print("Summary delta - unique cell counts", file=file)
        print(delta, file=file)
        print(file=file)

    # Genes removed, added
    for organism in census[CENSUS_DATA_NAME]:
        curr_genes = census[CENSUS_DATA_NAME][organism].ms["RNA"].var.read().concat().to_pandas()
        if organism not in previous_census[CENSUS_DATA_NAME]:
            # Nothing to diff against for a new organism; report and skip gene diffs
            print(
                f"New organism {organism} added with {len(curr_genes)} genes.",
                file=file,
            )
            print(file=file)
            continue

        prev_genes = previous_census[CENSUS_DATA_NAME][organism].ms["RNA"].var.read().concat().to_pandas()

        new_genes = set(curr_genes["feature_id"]) - set(prev_genes["feature_id"])
        if new_genes:
            print(f"Genes added in {organism} ({len(new_genes)}):", file=file)
            print(new_genes, file=file)
        else:
            print(f"No genes were added in {organism}.", file=file)
            print(file=file)

        removed_genes = set(prev_genes["feature_id"]) - set(curr_genes["feature_id"])
        if removed_genes:
            print(f"Genes removed in {organism} ({len(removed_genes)}):", file=file)
            print(removed_genes, file=file)
        else:
            print(f"No genes were removed in {organism}.", file=file)
            print(file=file)

    return 0


def create_args_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="cellxgene_census_summary")
    parser.add_argument("-c", "--census-version", default="latest", help="Version of the Census. Defaults to latest")
    subparsers = parser.add_subparsers(required=True, dest="subcommand")

    # BUILD
    subparsers.add_parser("summarize", help="Summarize the Census")

    # VALIDATE
    diff_parser = subparsers.add_parser("diff", help="Shows the diff with a previous Census version")
    diff_parser.add_argument("-p", "--previous-version", help="Version of the Census to diff")

    return parser


def main() -> int:
    parser = create_args_parser()
    args = parser.parse_args()
    assert args.subcommand in ["summarize", "diff"]

    if args.subcommand == "summarize":
        return display_summary(census_version=args.census_version)
    elif args.subcommand == "diff":
        return display_diff(census_version=args.census_version, previous_census_version=args.previous_version)

    return 0


if __name__ == "__main__":
    sys.exit(main())
