import os
import numpy as np
import pandas as pd
import yaml
import matplotlib.pyplot as plt
import matplotlib as mpl
import json
import argparse
import re
from src.signature import kmeans, nmf
from src.utils.vis import figure, subplots
from src.utils.signature import plot_signature_heatmap_full, plot_signature_heatmap_group

def main():
    # Load config
    parser = argparse.ArgumentParser() #hua added for differet configs
    parser.add_argument(
        "-c", "--config",
        help="Path(s) to YAML config."
    )
    args = parser.parse_args()

    # resolution order: CLI --config (in order) > env APP_CONFIG > ./config.yaml
    config_path = args.config
    config_name=re.sub(r'^(?:config_)?|\.ya?ml$', '',config_path)
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)

    signature_method = config['signature_config']['method']
    print(f'Using {signature_method} method for signature extraction')
    postprocess_folder = config['postprocess_config']['output']['path']+'/'+config_name #hua

    output_path = os.path.join(config['signature_config']['output']['path'], signature_method,config_name) #hua
    os.makedirs(output_path, exist_ok=True)

    data_types = ['upregulated', 'downregulated', 'conserved']
    for data_type in data_types:
        print(f"Processing {data_type} data")

        data_type_folder = os.path.join(postprocess_folder, data_type)
        metadata_path = os.path.join(data_type_folder, 'metadata.json')
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)

        data_save_path = os.path.join(output_path, data_type)
        os.makedirs(data_save_path, exist_ok=True)

        # Find loci/CAP group signatures using method of choice
        plot_dict = {}

        # Kmeans
        if signature_method == 'kmeans':
            diff_mat_path = os.path.join(data_type_folder, f'differential.npy')
            diff_mat = np.load(diff_mat_path)
            plot_dict = kmeans.get_signature(
                                    data_save_path=data_save_path, 
                                    diff_mat=diff_mat, 
                                    loci_list=metadata['filtered_gene_names'], 
                                    cap_list=metadata['filtered_cap_names'], 
                                    direction=data_type, 
                                    normalize=True,
                                    rerun_cluster=True
                                )
            plot_signature_heatmap_full(plot_dict, block_size=400, vmin=-10, vmax=10, show_ticks=True)
            plot_signature_heatmap_group(plot_dict)

        # NMF
        elif signature_method == 'nmf':
            # Process each celltype separately
            celltypes = [metadata['control_name'], metadata['experiment_name']]
            for celltype in celltypes:
                print(f'Processing celltype: {celltype}')
                celltype_mat_path = os.path.join(data_type_folder, f'{celltype}.npy')
                celltype_mat = np.load(celltype_mat_path)
            
                celltype_save_path = os.path.join(data_save_path, celltype)
                os.makedirs(celltype_save_path, exist_ok=True)

                plot_dict[celltype] = nmf.get_signature(
                                    data_save_path=celltype_save_path,  
                                    data_mat=celltype_mat,
                                    loci_list=metadata['filtered_gene_names'],
                                    cap_list=metadata['filtered_cap_names'],
                                    direction=data_type,
                                    rerun_cluster=True
                                )
                plot_signature_heatmap_full(plot_dict[celltype], block_size=400, vmin=-10, vmax=10, show_ticks=True)
                plot_signature_heatmap_group(plot_dict[celltype])

            # Compute similarity between celltype groups
            nmf.compute_group_similarity(plot_dict[celltypes[0]], plot_dict[celltypes[1]], celltypes, data_save_path)

if __name__ == "__main__":
    main()