"""
Utility script to update test_config.py with the latest resource sizes and shared gene counts.
Run this script whenever resource files are updated.

Usage:
    python tests/update_resource_sizes.py
"""

import os

from lineage import Lineage
from lineage.resources import Resources
from tests import simulate_snps

# Get the directory where this script is located
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))


class ResourceUpdater:
    def __init__(self):
        self.l = Lineage()

    def update_resource_sizes(self):
        resources = Resources()

        # Get resource sizes
        resource_sizes = {
            "knownGene_hg19": len(resources.get_knownGene_hg19()),
            "kgXref_hg19": len(resources.get_kgXref_hg19()),
            "cytoBand_hg19": len(resources.get_cytoBand_hg19()),
        }

        # Compute shared gene counts for each test case
        shared_gene_counts = {}

        # Default (HapMap2) test case (mimics test_find_shared_dna_two_chrom_shared)
        ind1 = simulate_snps(
            self.l.create_individual("ind1"),
            chrom="1",
            pos_max=111700001,
            pos_step=10000,
            genotype="AA",
        )
        ind2 = simulate_snps(
            self.l.create_individual("ind2"),
            chrom="1",
            pos_max=111700001,
            pos_step=10000,
            genotype="AA",
        )
        d = self.l.find_shared_dna(
            [ind1, ind2], shared_genes=True, genetic_map="HapMap2"
        )
        shared_gene_counts["default"] = {
            "len1": len(d["one_chrom_shared_genes"]),
            "len2": len(d["two_chrom_shared_genes"]),
        }

        # 1000G test case (mimics test_find_shared_dna_two_chrom_shared_1000G)
        ind1 = simulate_snps(
            self.l.create_individual("ind1"),
            chrom="1",
            pos_max=43800002,
            pos_step=10000,
            genotype="AA",
        )
        ind2 = simulate_snps(
            self.l.create_individual("ind2"),
            chrom="1",
            pos_max=43800002,
            pos_step=10000,
            genotype="AA",
        )
        d = self.l.find_shared_dna([ind1, ind2], shared_genes=True, genetic_map="CEU")
        shared_gene_counts["1000G"] = {"len1": len(d["one_chrom_shared_genes"])}

        # X_chrom_male test case (mimics test_find_shared_dna_X_chrom_two_individuals_male)
        ind1 = simulate_snps(
            self.l.create_individual("ind1"),
            chrom="X",
            pos_max=155270560,
            pos_step=1000,
            genotype="AA",
        )
        ind2 = simulate_snps(
            self.l.create_individual("ind2"),
            chrom="X",
            pos_max=155270560,
            pos_step=1000,
            genotype="AA",
        )
        d = self.l.find_shared_dna([ind1, ind2], shared_genes=True)
        shared_gene_counts["X_chrom_male"] = {
            "len1": len(d["two_chrom_shared_genes"]),
            "len2": len(d["one_chrom_shared_genes"]),
        }

        # X_chrom_female test case (mimics test_find_shared_dna_X_chrom_two_individuals_female)
        ind1 = simulate_snps(
            self.l.create_individual("ind1"),
            chrom="X",
            pos_max=155270560,
            pos_step=1000,
            genotype="AC",
        )
        ind2 = simulate_snps(
            self.l.create_individual("ind2"),
            chrom="X",
            pos_max=155270560,
            pos_step=1000,
            genotype="AC",
        )
        d = self.l.find_shared_dna([ind1, ind2], shared_genes=True)
        shared_gene_counts["X_chrom_female"] = {
            "len1": len(d["one_chrom_shared_genes"]),
            "len2": len(d["two_chrom_shared_genes"]),
        }

        # Write to test_config.py in the same directory as this script
        config_path = os.path.join(SCRIPT_DIR, "test_config.py")
        with open(config_path, "w") as f:
            f.write(
                '"""\nConfiguration file for test resource sizes and shared gene counts.\n'
                "This file is automatically generated by running update_resource_sizes.py.\n"
                'Do not edit manually; rerun update_resource_sizes.py to update values.\n"""\n\n'
            )
            f.write("# Resource sizes for hg19 data\n")
            f.write(f"RESOURCE_SIZES = {resource_sizes}\n\n")
            f.write("# Shared gene counts for specific test cases\n")
            f.write(f"SHARED_GENE_COUNTS = {shared_gene_counts}\n")

        print(f"Updated {config_path}")


if __name__ == "__main__":
    updater = ResourceUpdater()
    updater.update_resource_sizes()
