import argparse


def main(args):
    import glob
    import itertools
    import json
    import random

    import numpy as np

    with open(args.input_path, "r") as json_file:
        json_list = list(json_file)

    fixed_list = [
        [int(item) for item in one.split()] for one in args.position_list.split(",")
    ]
    global_designed_chain_list = [str(item) for item in args.chain_list.split()]
    my_dict = {}

    if not args.specify_non_fixed:
        for json_str in json_list:
            result = json.loads(json_str)
            all_chain_list = [
                item[-1:] for item in list(result) if item[:9] == "seq_chain"
            ]
            fixed_position_dict = {}
            for i, chain in enumerate(global_designed_chain_list):
                fixed_position_dict[chain] = fixed_list[i]
            for chain in all_chain_list:
                if chain not in global_designed_chain_list:
                    fixed_position_dict[chain] = []
            my_dict[result["name"]] = fixed_position_dict
    else:
        for json_str in json_list:
            result = json.loads(json_str)
            all_chain_list = [
                item[-1:] for item in list(result) if item[:9] == "seq_chain"
            ]
            fixed_position_dict = {}
            for chain in all_chain_list:
                seq_length = len(result[f"seq_chain_{chain}"])
                all_residue_list = (np.arange(seq_length) + 1).tolist()
                if chain not in global_designed_chain_list:
                    fixed_position_dict[chain] = all_residue_list
                else:
                    idx = np.argwhere(np.array(global_designed_chain_list) == chain)[0][
                        0
                    ]
                    fixed_position_dict[chain] = list(
                        set(all_residue_list) - set(fixed_list[idx])
                    )
            my_dict[result["name"]] = fixed_position_dict

    with open(args.output_path, "w") as f:
        f.write(json.dumps(my_dict) + "\n")

    # e.g. output
    # {"5TTA": {"A": [1, 2, 3, 7, 8, 9, 22, 25, 33], "B": []}, "3LIS": {"A": [], "B": []}}


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    argparser.add_argument("--input_path", type=str, help="Path to the parsed PDBs")
    argparser.add_argument(
        "--output_path", type=str, help="Path to the output dictionary"
    )
    argparser.add_argument(
        "--chain_list",
        type=str,
        default="",
        help="List of the chains that need to be fixed",
    )
    argparser.add_argument(
        "--position_list",
        type=str,
        default="",
        help="Position lists, e.g. 11 12 14 18, 1 2 3 4 for first chain and the second chain",
    )
    argparser.add_argument(
        "--specify_non_fixed",
        action="store_true",
        default=False,
        help="Allows specifying just residues that need to be designed (default: false)",
    )

    args = argparser.parse_args()
    main(args)