"""
Copyright 2018 Michael Yourshaw. All rights reserved.

myourshaw@gmail.com

Licensed under the MIT license.

Adapted from gsa.reshape.concordance.table.R
<https://github.com/broadinstitute/gsalib/blob/master/R/gsa.reshape.concordance.table.R>

Given a GATKReport generated by GenotypeConcordance (as output by \code{gsa.read.gatkreport}),
this function reshapes the concordance for a specified sample into a matrix with the
EvalGenotypes in rows and the CompGenotypes in columns.
(see the documentation for GenotypeConcordance for the definition of Eval and Comp)

Output looks like:
CompGenotypes   NO_CALL  HOM_REF      HET  HOM_VAR  UNAVAILABLE  MIXED
EvalGenotypes
      NO_CALL         0        0        0        0            0      0
      HOM_REF         0        0        0        0            0      0
          HET         0        0    13463       90         3901      0
      HOM_VAR         0        0     2935    18144         4448      0
  UNAVAILABLE         0        0  2053693  1326112        11290      0
        MIXED         0        0        0        0            0      0
"""

import pandas as pd
from gsalib import GatkReport


def reshape_concordance_table(report, table_name="GenotypeConcordance_Counts", sample_name="ALL"):
    """
    Given a GATKReport generated by GenotypeConcordance (as output by \code{gsa.read.gatkreport}),
    this function reshapes the concordance for a specified sample into a matrix with the
    EvalGenotypes in rows and the CompGenotypes in columns.
    (see the documentation for GenotypeConcordance for the definition of Eval and Comp)
    :param report: a GatkReport object
    :param table_name: the name of a table in the report
    :param sample_name: the name of a sample in the table (or  'ALL')
    :return: a DataFrame of eval/comp counts, indexed by EvalGenotypes, and columns of CompGenotypes
    """
    if table_name is not None:
        data = report[table_name]
    else:
        return None

    # the genotypes of interest, based on the table's column names,
    # which have the form eval_genotype_comp_genotype
    # ignoring columns Sample and Mismatching_Alleles, and maybe others
    possible_genotypes = ['NO_CALL', 'HOM_REF', 'HET', 'HOM_VAR', 'UNAVAILABLE', 'MIXED']

    # transpose the table into a 1-column DataFrame indexed by eval_genotype_comp_genotype strings
    d = data.loc[data['Sample'] == sample_name].T

    # create a list of the genotype counts ordered by eval, comp
    m = []
    for eval_genotype in possible_genotypes:
        eval_df = d[d.index.str.startswith(eval_genotype)]
        for comp_genotype in possible_genotypes:
            comp_df = eval_df[eval_df.index.str.endswith(comp_genotype)]
            m.append(comp_df.values[0][0])

    # create a matrix-like Series with comp
    s = pd.Series(m).values.reshape(len(possible_genotypes), len(possible_genotypes))
    df = pd.DataFrame(s, columns=possible_genotypes, index=possible_genotypes).rename_axis(
        'EvalGenotypes').rename_axis("CompGenotypes", axis="columns")

    return df


import os
this_dir = os.path.dirname(os.path.realpath(__file__))
report = GatkReport(os.path.join(this_dir, 'Concordance_NA12878.lValExome0030_Platinum.OMIM.summary.grp'))
df = reshape_concordance_table(report)
print(df)
