from collections import Counter
import gzip
import json

from eva_cttv_pipeline.trait_mapping.trait import Trait


# TODO: unify the acceptable clinical significance levels project-wide
ACCEPTABLE_CLINICAL_SIGNIFICANCE = ['pathogenic', 'likely pathogenic', 'protective', 'association', 'risk_factor',
                                    'affects', 'drug response']


def parse_trait_names(filepath: str) -> list:
    """
    For a file containing ClinVar records in the TSV format, return a list of Traits for the records in the file. Each
    Trait object contains trait name, how many times it occurs in the input file, and whether it is linked to an NT
    expansion variant.

    Trait occurrence count is calculated based on all unique (AlleleID, RCV, trait name) tuples in the input file. This
    is because each such tuple will, generally speaking, correspond to one output evidence string. So if we want to
    gauge which trait names are more important to curate, we need to consider how many such tuples it appears in.

    The reason we need to keep track of only *unique* tuples is because some (most) alleles will appear twice in the
    document with coordinates for GRCh37 and GRCh38, and we don't want to count them twice.

    Traits which are implicated in "NT expansion" variants are marked using a special field, because their curation is
    of highest importance even if the number of records which they are linked to is low.

    :param filepath: Path to a gzipped file containing ClinVar TSV summary.
    :return: A list of Trait objects.
    """

    # Tracks unique (AlleleID, RCV, trait name) tuples
    unique_association_tuples = set()

    # Tracks all traits which are at least once implicated in "NT expansion", or nucleotide repeat expansion, variants.
    # Their curation is of highest importantce regardless of how many records they are actually associated with.
    nt_expansion_traits = set()

    with gzip.open(filepath, "rt") as clinvar_summary:
        header = clinvar_summary.readline().rstrip().split('\t')
        for line in clinvar_summary:
            values = line.rstrip().split('\t')
            data = dict(zip(header, values))

            # Check if the record should be processed given its level of clinical significance
            acceptable_clinical_significance_present = False
            for clinical_significance in data['ClinicalSignificance'].split(','):
                if clinical_significance.strip().lower() in ACCEPTABLE_CLINICAL_SIGNIFICANCE:
                    acceptable_clinical_significance_present = True
            if not acceptable_clinical_significance_present:
                continue

            # Extract relevant fields
            is_nt_expansion_variant = data['Type'] == 'NT expansion'
            allele_id = data['#AlleleID']
            traits = set(data['PhenotypeList'].split(';'))
            rcv_ids = set(data['RCVaccession'].split(';'))

            # Process all (trait, rcv) records
            for trait, rcv_id in zip(traits, rcv_ids):
                unique_association_tuples |= {(trait, rcv_id, allele_id)}
                if is_nt_expansion_variant:
                    nt_expansion_traits |= {trait}

    # Count trait occurrences
    trait_names = [t[0] for t in unique_association_tuples]
    traits = []
    for trait_name, trait_frequency in Counter(trait_names).items():
        associated_with_nt_expansion = trait_name in nt_expansion_traits
        traits.append(Trait(name=trait_name.lower(), frequency=trait_frequency,
                            associated_with_nt_expansion=associated_with_nt_expansion))

    return traits

@@ -3074,21 +3074,21 @@
 n_tuples
- %7C= %7B
+.add(
 (trait, 
@@ -3105,17 +3105,17 @@
 lele_id)
-%7D
+)
 %0A       
@@ -3194,19 +3194,19 @@
 aits
- %7C= %7B
+.add(
 trait
-%7D
+)
 %0A%0A  

Apply suggestions from code review