import modal
import os

MODEL_STORE_VOLUME_DIR = "/volume"
CLUSTER_DIR = "/app2/suno/modal"
CLUSTER_SSD_DIR = "/mnt/localdisk"


def is_modal():
    return os.environ.get("MODAL_REGION") is not None


MODEL_STORE_VOLUME_PREFIX = MODEL_STORE_VOLUME_DIR + "/models"
if is_modal():
    print("Running on Modal")
    MODEL_STORE_VOLUME_PREFIX = MODEL_STORE_VOLUME_DIR + "/models"
else:
    print("Running on non-Modal")

    def compare_directories_by_name(dir1, dir2):
        """
        Compare directories based only on file names.
        Returns the percentage of matching files.
        """
        # Get all file names recursively (including in subdirectories)
        files1 = set()
        files2 = set()

        # Walk through first directory
        for root, _, files in os.walk(dir1):
            rel_path = os.path.relpath(root, dir1)
            for f in files:
                # Store relative paths to compare directory structure too
                rel_file = os.path.join(rel_path, f).replace("\\", "/")
                if rel_file.startswith("./"):
                    rel_file = rel_file[2:]
                files1.add(rel_file)

        # Walk through second directory
        for root, _, files in os.walk(dir2):
            rel_path = os.path.relpath(root, dir2)
            for f in files:
                rel_file = os.path.join(rel_path, f).replace("\\", "/")
                if rel_file.startswith("./"):
                    rel_file = rel_file[2:]
                files2.add(rel_file)

        # Calculate intersection and union
        common_files = files1.intersection(files2)
        all_files = files1.union(files2)

        # Calculate percentage
        if len(all_files) == 0:
            return 0.0  # Empty directories

        match_percentage = (len(common_files) / len(all_files)) * 100

        # Print summary
        print(f"Dir1: {len(files1)} files")
        print(f"Dir2: {len(files2)} files")
        print(f"Common files: {len(common_files)}")
        print(f"All unique files: {len(all_files)}")
        print(f"Match percentage: {match_percentage:.2f}%")

        return match_percentage

    MODEL_STORE_VOLUME_PREFIX = CLUSTER_DIR + "/models"
    CLUSTER_SSD_PREFIX = CLUSTER_SSD_DIR + "/models"
    similarity = compare_directories_by_name(MODEL_STORE_VOLUME_PREFIX, CLUSTER_SSD_PREFIX)
    print(f"Similarity: {similarity:.2f}%")
    if similarity == 100:
        MODEL_STORE_VOLUME_PREFIX = CLUSTER_SSD_PREFIX
    else:
        print(
            "Directories are not identical. Use disk model instead, if you want to use ssd models. run python -m suno_utils.worker.modal_manager_model_store sync"
        )


model_store_volume = modal.Volume.from_name("model-store", create_if_missing=False)

GPT_PATH_DICT = {
    "7b": "/georg/checkpoints/chirp_v2_5/7b_ft_v3.pt",
    "7b_dpo": "/tony/fine_tune_7b/dpo_v0_t1.pt",
    "7b_ipo": "/tony/fine_tune_7b/dpo_7b_v0_t12.pt",
    "7b_special": "/tony/fine_tune_7b/dpo_7b_v1_t5.pt",
    "7b_fast": "/tony/tmp/model_7b_s5_t1_v18.pt",
    "v2": "/georg/checkpoints/chirp_v2/xxl.pt",
    "13b": "/georg/tmp/model_13b.pt",
    "13b_ft": "/tony/tmp/model_13b_2h.pt",  # 2h ft with infilling, etc
    "13b_ft_1": "/tony/dpo_132b/dpo_13b_2h_v3.pt",  # dpo-ed 13b 2h
    "13b_special_2": "/tony/dpo_132b/dpo_13b_v6.pt",
    "13b_special_8": "/tony/dpo_132b/dpo_13b_v13.pt",  # curent production model
    "13b_short": "/tony/dpo_132b/dpo_13b_v13.pt",  # s8-v3 dpo-ed
    "13b_special_31": "/tony/dpo_132b/dpo_13b_s29_v6.pt",  # s29-6 dpo-ed
    "13b_special_32": "/tony/dpo_132b/dpo_13b_s31_v9.pt",  # s31-9 dpo-ed
    "13b_tech": "/tony/dpo_132b/dpo_13b_s31_v9.pt",  # same as production, for eng testing
    "13b_special_test": "/tony/dpo_132b/dpo_13b_s31_v9.pt",  # s32-1 dpo-ed person
    "13b_special_test_2": "/tony/dpo_132b/dpo_13b_s31_v9.pt",  # s32 person negative
    "13b_special_test_3": "/tony/dpo_132b/dpo_13b_s32_v30_person_1_1k.pt",  # s32 person positive
    "13b_special_test_diff": "/tony/dpo_132b/dpo_13b_s31_v9.pt",  # s8-v3 dpo-ed person
    "13b_special_33": "/tony/dpo_132b/dpo_13b_s32_v29.pt",  # s32-29 dpo-ed person
    "13b_special_32_fast": "/tony/dpo_132b/dpo_13b_s31_v9.pt",  # s31-9 dpo-ed
    "13b_upload": "/tony/dpo_132b/dpo_13b_v6.pt",  # upload model
    "13b_upload_4": "/tony/dpo_132b/dpo_13b_upload_v4.pt",  # test upload model
    "30b": "/tony/tmp/model_30b_fix_ft2_20k.pt",  # 30b ft-ed with infer token fix
    "30b_t2": "/tony/tmp/model_30b_ft2_t1_v12.pt",  # 30b 1st dpo-ed (v12)
    "30b_t3": "/tony/tmp/model_30b_ft2_t2_v10.pt",  # 30b 2nd dpo-ed
    "30b_t5": "/tony/tmp/model_30b_t4_v25_beta20.pt",  # 30b 3rd dpo
    "30b_t6": "/tony/tmp/model_30b_t5_v15_b40.pt",  # 30b 4th dpo
    "30b_t6_infill": "/tony/tmp/model_30b_t5_v15_b40.pt",  # 30b 4th dpo
    "30b_t6_test": "/tony/dpo_30b/model_30b_t6_v35_infill.pt",  # 30b 4th dpo
    "30b_t6_tech": "/tony/tmp/model_30b_t5_v15_b40.pt",  # 30b 4th dpo, for eng testing
    "30b_t6_eval": "/sara/tony_30b_dpo_2_28.pt",
    "30b_test_diff": "/tony/tmp/model_30b_t5_v15_b40.pt",  # 30b 4th dpo
    "30b_classical": "/tony/tmp/model_30b_t3_classical_v1_5k.pt",  # 30b classical
    "30b_dance": "/tony/tmp/model_30b_t3_dg_v0.pt",  # 30b dance
    # v4.5
    "3b_sem_orig": "/tony/sem/model_45_3b_jan28_c_fix.pt",
    "3b_sem_test": "/tony/tmp/model_45_3b.pt",
    "6b_sem": "/tony/sem/model_45_6b_apr15_sft.pt",
    "6b_sem_eval": "/sara/bluejay_t1-r7_2025-05-21_22-39-00.pt",
    "6b_sem_test": "/tony/sem/model_45_6b_bjay_sft_0519_t15_r7_beta10_d67.pt",
    "6b_sem_test_2": "/tony/sem/model_45_6b_bjay_sft_0519_t1_r7.pt",
    "6b_sem_tech": "/tony/sem/model_45_6b_apr15_sft_dpo_v10_11.pt",  # borrow this for a bit
    "6b_sem_t1": "/tony/sem/model_45_6b_apr15_sft_dpo_v10_11.pt",
    "6b_sem_task": "/tony/sem/model_45_6b_apr15_sft_dpo_v10_11.pt",
    "6b_sem_bluejay": "/tony/sem/model_45_6b_bjay_sft_0519_t1_r9_d28.pt",  # aura g
    "6b_sem_t2": "/tony/sem/model_45_6b_bjay_sft_0519_t1_r7.pt",
    "6b_sem_bluejay_test": "/tony/sem/model_45_6b_bjay_sft_0519_t1_r7.pt",  # aura a
    "6b_sem_bluejay_test_2": "/tony/sem/model_45_6b_bjay_sft_0519_t8_r5_beta10_l_d60.pt",  # aura h
}

DIFFUSION_DEFAULT_V1_PATH = "/tony/tmp/diff/dit_v6_dpo_t11_9k_5e6_b100.pt"
DIFFUSION_DEFAULT_V2_PATH = "/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_t1_18_cs0.pt"

DIFFUSION_PATH_DICT = {
    "diff_v1": DIFFUSION_DEFAULT_V1_PATH,
    "diff_v1_test": "/tony/tmp/diff/dit_v6_dpo_t11_9k_5e6_b100.pt",
    "diff_v2": DIFFUSION_DEFAULT_V2_PATH,
    "diff_v2_test": "/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d4_v28.pt",
    "diff_v2_data": "/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt",
    "stems_v0": "/victor/checkpoints/diffusion/stems_10s_9369.pt",
    "diff_seeds_v0": "/christian/checkpoints/seeds/seeds_v0_step_60k.pt",
    # "stems_v1": "/victor/checkpoints/diffusion/stems_v1p2.pt",
    "stems_v1": "/victor/checkpoints/diffusion/stems_multi.pt",
    "stems_v1_8_output": "/victor/checkpoints/diffusion/stems_8_stem.pt",
    "stems_v1_12_output": "/victor/checkpoints/diffusion/stems_12_stem.pt",
    # "13b_special_test": "/tony/tmp/diff/dit_v1_dpo_t1_v7_1_3k.pt",
    # "13b_special_test_3": "/tony/tmp/diff/dit_v4_dpo_t7_3k_5e6_b100.pt",
    "13b_special_test_diff": "/tony/tmp/diff/dit_v6_dpo_t3_30k_5e6_b100.pt",
    "30b_test_diff": "/tony/tmp/diff/dit_v2_dpo_t2_v1_3k.pt",
    "30b_t6_infill": DIFFUSION_DEFAULT_V2_PATH,
    "3b_sem_orig": DIFFUSION_DEFAULT_V2_PATH,
    "3b_sem_test": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem_eval": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem_test": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem_test_2": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem_task": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem_t1": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem_tech": DIFFUSION_DEFAULT_V2_PATH,
    "6b_sem_bluejay": "/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt",
    "6b_sem_t2": "/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt",
    "6b_sem_bluejay_test": "/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt",
    "6b_sem_bluejay_test_2": "/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt",
    # "30b_t6_test": "/tony/tmp/diff/dit_v1_dpo_t1_v7_1_3k.pt",
}

CODEC_PATH_DICT = {
    # Codec checkpoint path correlated with diffusion versions.
    "v_vae_25_tuned_2": "/minz/models/dac_vae_tuned_25hz.pth",
    "v_vae_25_peaq_1": "/christian/25hz_vae_peaq_kl_0.005.pth",
}


# All other things that needs to be uploaded to Modal
GPT_CKPT_PATH = "/georg/checkpoints/chirp_v2_5/7b_ft.pt"
TOKENIZER_PATH = "/georg/trained_models/chirp_v2/tokenizer_60k.json"
SEMANTIC_CKPT_PATH = "/georg/models/semantic/mert_25.pt"
SEMANTIC_CENTROIDS_PATH = "/georg/models/semantic/mert_25_2x4k.npy"
CODEC_CKPT_PATH = "/georg/models/codec/dac_2c_25x12.pt"
HOOT_CKPT_PATH = "/checkpoints/hoot_v3/hoot_ckpt.pt"
HOOT_TOKENIZER_PATH = "/checkpoints/hoot_v3/tokenizer.model"
FASTTEXT_CKPT_PATH = "/georg/trained_models/chirp_v1/lid.176.bin"


TOKENIZER_FILEPATH = "/georg/models/tokenizers/tokenizer_60k.json"
SEMANTIC_MODEL_FILEPATH = "/georg/models/semantic/mert_25.pt"
SEMANTIC_CLUSTERS_FILEPATH = "/georg/models/semantic/mert_25_2x4k.npy"
CODEC_FILEPATH = "/christian/25hz_vae_peaq_kl_0.005.pth"
DIT_MODEL_FILEPATH = "/georg/tmp/2b_prefix_ft.pt"
