#!/bin/env python3
"""
Psi-calc is an algorithm for clustering protein multiple sequence alignments (MSAs)
utilizing normalized mutual information.

    Copyright (C) 2020 Thomas Townsley, MSSE, Joe Deweese, PhD., Kirk Durston, PhD.,
    Timothy Wallace, PhD, Salvador Cordova, PhD et. al.

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.

Contact: thomas@mandosoft.dev, joe.deweese@lipscomb.edu, kirkdurston@gmail.com
"""

import re
import time
import csv
import pandas as pd
import numpy as np
from itertools import combinations
from sklearn.metrics.cluster import normalized_mutual_info_score as nmis


def select_subset(c_list: list, s: int):
    """Selects the width of the sample subset from the MSA"""
    each_nth_col = c_list[::s]
    return each_nth_col


def durston_schema(df: pd.DataFrame, value: int) -> pd.DataFrame:
    """Labels index based on first value given. For example making the value 3
    will label the columns 3-N."""
    df.columns = range(len(df.columns))
    label_val = value
    df = df.rename(columns=lambda x: int(x) + label_val)

    return df


def deweese_schema(df: pd.DataFrame) -> pd.DataFrame:
    """Labels data based on the range on the first row of the MSA.
    For example, if the first row is labelled TOP2 YEAST/59-205, then all
    columns with individuals in the first row are kept and labeled 59-205."""

    try:
        df = df.rename(columns=lambda x: x - 1)
        first_row_ix = df.index[0]
        ix_label = first_row_ix.rsplit('/', 1)
        ix_label = ix_label[1]
        ix_label = ix_label.rsplit('-', 1)
        df_label = int(ix_label[0])
        column_lab_dict = dict()

        pattern = '^-'

        for i in df:
            if not re.search(pattern, df[i].iloc[0]):
                column_lab_dict[df.columns[i]] = df_label
                df_label += 1
            else:
                column_lab_dict[df.columns[i]] = ''
        df = df.rename(columns=column_lab_dict)
        df = df.drop(columns=[''])
        df = df.rename(columns=lambda x: int(x))
    except IndexError or KeyError:
        exit(1)
    return df


def read_txt_file_format(file) -> pd.DataFrame:
    """Reads FASTA files or text file-based MSAs into a dataframe."""
    nucs_dict = dict()
    with open(file, "r") as a_file:
        string_without_line_breaks = ""
        for line in a_file:
            stripped_line = line.rstrip()
            string_without_line_breaks += stripped_line
        a_file.close()

    vals = string_without_line_breaks.split('>')
    for line in vals:
        line = re.split('(\w+/\d*-?\d*)', line)
        line.remove('')
        for i in enumerate(line):
            g = list(line[1])
            nucs_dict.update({line[0]: g})

    df = pd.DataFrame.from_dict(nucs_dict, orient='index')
    df = df.replace({'.': '-'})
    df.index.name = 'SEQUENCE_ID'
    df.drop_duplicates(inplace=True)

    return df


def read_csv_file_format(file) -> pd.DataFrame:
    """Reads CSV MSAs into a dataframe."""
    df = pd.read_csv(file, encoding='utf-8', engine='c', header=None)
    df = df.rename(columns={df.columns[0]: 'SEQUENCE_ID'})
    df = df.set_index('SEQUENCE_ID', drop=True)
    df.drop_duplicates(inplace=True)

    return df


def return_sr_mode(msa: np.ndarray, m_map: dict, c: list, csv_dict: dict, list_store: list, k) -> tuple:
    """Calculates the sr_mode and new mode of a cluster and returns both."""
    cc = len(list(combinations(c, 2)))
    if len(c) < 2:
        sr_mode, new_mode = None, None
        return sr_mode, new_mode
    elif len(c) == 2:
        i, j = m_map.get(c[0]), m_map.get(c[1])
        max_sum = nmis(msa[:, i], msa[:, j], average_method='geometric')
        new_mode = c
        if max_sum > 1.0:
            print("\nERROR: Score returned greater than 1.0: ")
            print(c, max_sum)
            exit(1)
    else:
        A = [[column, []] for column in range(len(c))]
        D = len(A)
        shift = 0
        for loc, i in enumerate(c):
            for location, j in enumerate(c):
                if location != loc and location > shift:
                    l, r = m_map.get(i), m_map.get(j)
                    A[loc][1].append(nmis(msa[:, l], msa[:, r], average_method='geometric'))
            t = loc
            q = loc + 1
            if q != D:
                while q != D:
                    A[q][1].append(A[loc][1][t])
                    q += 1
                    t += 1
                shift += 1

        mode_map = np.array([np.sum(v) for (k, v) in A], dtype='d')
        mode_loc = int(np.argmax(mode_map))
        max_sum = np.amax(mode_map)
        new_mode = return_new_mode(mode_loc, c)

    sr_mode = max_sum / cc
    if k == "pairwise" or k == "post-agg":
        list_store.append([sr_mode, new_mode])
    if k != "pairwise":
        csv_dict[tuple(sorted(tuple(c)))] = [round(sr_mode, 6), k]

    return sr_mode, new_mode


def return_new_mode(location: int, c: list) -> list:
    """Calculates the new mode of a cluster and returns the updated cluster.
    The mode of a cluster is the left-most column in the dataframe."""
    mode = c[location]
    c.remove(mode)
    c.append(mode)
    c.reverse()
    return c


def list_contains(list1, list2, capture_sets):
    """Finds clusters with overlapping attributes"""
    set1 = set(list1)
    for each in list2:
        set2 = set(each)
        if set1 != set2 and set1.intersection(set2):
            set1 = set1.union(set2)
        else:
            pass
    if set1 not in capture_sets:
        capture_sets.append(set1)


def write_output_data(spread: int, csv_dict: dict):
    """Writes the CSV output file"""
    filename = "data_out_width" + str(spread) + ".csv"
    a_file = open(filename, "w")
    writer = csv.writer(a_file)
    writer.writerow(["Cluster", "Sr_mode", "Discovered"])
    for key, value in csv_dict.items():
        val1, val2 = value
        writer.writerow([key, val1, val2])
    a_file.close()
    return filename


def get_unique_elements(df: pd.DataFrame) -> np.ndarray:
    """Returns all unique elements found in a multiple
    sequence alignment."""
    U = np.array([])
    for name, seq in df.iteritems():
        U = np.append(U, seq.unique())

    return np.unique(U)


def rectify_sequences(df: pd.DataFrame) -> pd.DataFrame:
    """Corrects symbols in the multiple sequence alignment."""
    alpha = get_unique_elements(df)
    pattern = '[-#?.]'
    for symbol in alpha:
        if symbol.isupper() and len(symbol) == 1:
            pass
        else:
            if symbol.islower() and len(symbol) == 1:
                df.replace(symbol, symbol.upper(), inplace=True)
            elif re.search(pattern, symbol):
                df.replace(symbol, '-', inplace=True)
            elif len(symbol) > 1:
                df.replace(symbol, '-', inplace=True)
            else:
                print("FAILED: Bad symbol in MSA: ", symbol)
                exit(1)
    print("\nBad symbols were discovered and replaced. New symbols:\n", get_unique_elements(df))
    return df


def encode_msa(df: pd.DataFrame) -> np.ndarray:
    """Naive integer encoding. Encodes amino acid values with
    unique integers and returns a numpy matrix"""
    ent = 0
    alpha = get_unique_elements(df)
    if alpha.size > 22:
        df = rectify_sequences(df)
        alpha = get_unique_elements(df)
    for symbol in alpha:
        if symbol == '-':
            df.replace(symbol, 0, inplace=True)
        else:
            ent += 1
            df.replace(symbol, ent, inplace=True)
    if ent > 22:
        print("FAILED: Too many labels encoded: ", ent)
        exit(1)
    print("\nNumber of labels encoded: ", (ent + 1), " including gaps.")
    np_matrix = df.to_numpy()

    return np_matrix


def find_clusters(spread: int, df: pd.DataFrame) -> dict:
    """Discovers cluster sites with high shared normalized mutual information.
    Provide a dataframe and a sample spread-width. Returns a dictionary."""
    print("\nEncoding MSA. This may take a while.")
    k = "pairwise"

    csv_dict = dict()
    start_time = time.time()
    hash_list = list()
    msa_index = df.columns.tolist()
    msa_map = {k: v for v, k in enumerate(msa_index)}
    num_msa = encode_msa(df)
    subset = select_subset(msa_index, spread)
    subset_list = [[z] for z in subset]

    print("\nLooking for strong pairwise clusters...")
    for item, each in enumerate(subset_list):
        max_rii, best_cluster = 0, None
        for location, cluster in enumerate(msa_index):
            cluster_mode = msa_map.get(cluster)
            subset_mode = msa_map.get(each[0])
            if subset_mode != cluster_mode:
                rii = nmis(num_msa[:, subset_mode], num_msa[:, cluster_mode], average_method='geometric')
                if rii > max_rii:
                    max_rii, best_cluster = rii, location
        subset_list[item].append(msa_index[best_cluster])
        print("pair located: ", subset_list[item])

    for cluster in subset_list:
        return_sr_mode(num_msa, msa_map, cluster, csv_dict, hash_list, k)
    sorted_list = sorted(hash_list, key=lambda x: x[0], reverse=True)
    N = int(len(sorted_list) * 1 / 3)
    out_list = sorted_list[:N]
    dataframe_label_list = [j for x, j in out_list]

    # Check for repeat attributes between pairwise clusters
    capture_list = list()
    for each in dataframe_label_list:
        list_contains(each, dataframe_label_list, capture_list)
    final_df_set = [list(sets) for sets in capture_list]

    unranked = list()
    k = "post-agg"  # means clusters which were merged during post-aggregation
    for cluster in final_df_set:
        return_sr_mode(num_msa, msa_map, cluster, csv_dict, unranked, k)

    C = sorted(unranked, key=lambda x: x[0], reverse=True)
    for x, j in C:
        for col in j:
            msa_index.remove(col)
    print("\nTop ranked clusters:")
    num = 0
    for each in C:
        num += 1
        print(num, ": ", each[1])
    # Sets the number of clusters we're actually iterating over
    R = len(C)
    for remaining in msa_index:
        C.append([0, [remaining]])

    # Stage Two: Move through the pairs in the list and find their best attribute
    k = len(C)
    num_clusters = len(C[0:R])

    while len(C) >= 2:
        print("\nNumber of Clusters Remaining: ", num_clusters)
        i = 0
        while i < num_clusters:
            location = None
            cluster_mode = msa_map.get(C[i][1][0])
            max_rii, best_cluster = 0, None
            for loc, entry in enumerate(C):
                sr_mode, cluster = entry
                attr_mode = msa_map.get(cluster[0])
                if cluster_mode != attr_mode:
                    rii = nmis(num_msa[:, attr_mode], num_msa[:, cluster_mode], average_method='geometric')
                    if rii > max_rii:
                        max_rii, best_cluster, location = rii, cluster, loc
            C[i][1] = C[i][1] + best_cluster
            C[i][0], C[i][1] = return_sr_mode(num_msa, msa_map, C[i][1], csv_dict, 0, k)
            C.pop(location)
            i += 1
            k -= 1
            print("k = ", k)
            if len(C) == 2:
                break
            if location <= num_clusters:
                num_clusters -= 1

        # Re-sort the list
        C = sorted(C, key=lambda x: x[0], reverse=True)
        print(" --> Next run at ", len(C))

        if num_clusters <= 2:
            break
    print("\n\n--- took " + str(time.time() - start_time) + " seconds ---")

    return csv_dict
