import glob
import json
import random

import numpy as np
import pandas as pd


def softmax(x, T):
    return np.exp(x / T) / np.sum(np.exp(x / T), -1, keepdims=True)


def parse_pssm(path):
    data = pd.read_csv(path, skiprows=2)
    floats_list_list = []
    for i in range(data.values.shape[0]):
        str1 = data.values[i][0][4:]
        floats_list = []
        for item in str1.split():
            floats_list.append(float(item))
        floats_list_list.append(floats_list)
    np_lines = np.array(floats_list_list)
    return np_lines


np_lines = parse_pssm(
    "/home/swang523/RLcage/capsid/monomersfordesign/8-16-21/pssm_rainity_final_8-16-21_int/build_0.2089_0.98_0.4653_19_2.00_0.005745.pssm"
)

mpnn_alphabet = "ACDEFGHIKLMNPQRSTVWYX"
input_alphabet = "ARNDCQEGHILKMFPSTWYV"

permutation_matrix = np.zeros([20, 21])
for i in range(20):
    letter1 = input_alphabet[i]
    for j in range(21):
        letter2 = mpnn_alphabet[j]
        if letter1 == letter2:
            permutation_matrix[i, j] = 1.0

pssm_log_odds = np_lines[:, :20] @ permutation_matrix
pssm_probs = np_lines[:, 20:40] @ permutation_matrix

X_mask = np.concatenate([np.zeros([1, 20]), np.ones([1, 1])], -1)


def softmax(x, T):
    return np.exp(x / T) / np.sum(np.exp(x / T), -1, keepdims=True)


# Load parsed PDBs:
with open("/home/justas/projects/cages/parsed/test.jsonl", "r") as json_file:
    json_list = list(json_file)

my_dict = {}
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"]
    pssm_dict = {}
    for chain in all_chain_list:
        pssm_dict[chain] = {}
        pssm_dict[chain]["pssm_coef"] = (
            np.ones(len(result["seq_chain_A"]))
        ).tolist()  # a number between 0.0 and 1.0 specifying how much attention put to PSSM, can be adjusted later as a flag
        pssm_dict[chain]["pssm_bias"] = (
            softmax(pssm_log_odds - X_mask * 1e8, 1.0)
        ).tolist()  # PSSM like, [length, 21] such that sum over the last dimension adds up to 1.0
        pssm_dict[chain]["pssm_log_odds"] = (pssm_log_odds).tolist()
    my_dict[result["name"]] = pssm_dict

# Write output to:
with open("/home/justas/projects/lab_github/mpnn/data/pssm_dict.jsonl", "w") as f:
    f.write(json.dumps(my_dict) + "\n")