# -*- coding: utf-8 -*-
# stdlib modules
import logging
import re
import pathlib
import tempfile

from datetime import datetime
from pprint import pprint as pp

# Third party modules
import pymongo
from pymongo.errors import (DuplicateKeyError, BulkWriteError)

from cyvcf2 import VCF
from intervaltree import IntervalTree

# Local modules
from scout.parse.variant.headers import (parse_rank_results_header,
                                         parse_vep_header)
from scout.parse.variant.rank_score import parse_rank_score

from scout.parse.variant import parse_variant
from scout.build import build_variant

from scout.exceptions import IntegrityError

from scout.constants import (CHROMOSOMES, FILE_TYPE_MAP)

LOG = logging.getLogger(__name__)

class VariantLoader(object):

    """Methods to handle variant loading in the mongo adapter"""

    def update_variant(self, variant_obj):
        """Update one variant document in the database.

        This means that the variant in the database will be replaced by variant_obj.

        Args:
            variant_obj(dict)

        Returns:
            new_variant(dict)
        """
        LOG.debug('Updating variant %s', variant_obj.get('simple_id'))

        new_variant = self.variant_collection.find_one_and_replace(
            {'_id': variant_obj['_id']},
            variant_obj,
            return_document=pymongo.ReturnDocument.AFTER
        )
        return new_variant

    def update_variant_rank(self, case_obj, variant_type='clinical', category='snv'):
        """Updates the manual rank for all variants in a case

        Add a variant rank based on the rank score
        Whenever variants are added or removed from a case we need to update the variant rank

        Args:
            case_obj(Case)
            variant_type(str)
        """
        # Get all variants sorted by rank score
        variants = self.variant_collection.find({
            'case_id': case_obj['_id'],
            'category': category,
            'variant_type': variant_type,
        }).sort('rank_score', pymongo.DESCENDING)

        LOG.info("Updating variant_rank for all variants")

        requests = []

        for index, var_obj in enumerate(variants):
            if len(requests) > 5000:
                try:
                    self.variant_collection.bulk_write(requests, ordered=False)
                    requests = []
                except BulkWriteError as err:
                    LOG.warning("Updating variant rank failed")
                    raise err

            operation = pymongo.UpdateOne(
                {'_id': var_obj['_id']},
                {
                    '$set': {
                        'variant_rank': index + 1,
                    }
                })
            requests.append(operation)

        #Update the final bulk
        try:
            self.variant_collection.bulk_write(requests, ordered=False)
        except BulkWriteError as err:
            LOG.warning("Updating variant rank failed")
            raise err

        LOG.info("Updating variant_rank done")

    def update_variant_compounds(self, variant, variant_objs = None):
        """Update compounds for a variant.

        This will add all the necessary information of a variant on a compound object.

        Args:
            variant(scout.models.Variant)
            variant_objs(dict): A dictionary with _ids as keys and variant objs as values.

        Returns:
            compound_objs(list(dict)): A dictionary with updated compound objects.

        """
        compound_objs = []
        for compound in variant.get('compounds', []):
            not_loaded = True
            gene_objs = []
            # Check if the compound variant exists
            if variant_objs:
                variant_obj = variant_objs.get(compound['variant'])
            else:
                variant_obj = self.variant_collection.find_one({'_id': compound['variant']})
            if variant_obj:
                # If the variant exosts we try to collect as much info as possible
                not_loaded = False
                compound['rank_score'] = variant_obj['rank_score']
                for gene in variant_obj.get('genes', []):
                    gene_obj = {
                        'hgnc_id': gene['hgnc_id'],
                        'hgnc_symbol': gene.get('hgnc_symbol'),
                        'region_annotation': gene.get('region_annotation'),
                        'functional_annotation': gene.get('functional_annotation'),
                    }
                    gene_objs.append(gene_obj)
                    compound['genes'] = gene_objs

            compound['not_loaded'] = not_loaded
            compound_objs.append(compound)

        return compound_objs

    def update_compounds(self, variants):
        """Update the compounds for a set of variants.

        Args:
            variants(dict): A dictionary with _ids as keys and variant objs as values

        """
        LOG.debug("Updating compound objects")

        for var_id in variants:
            variant_obj = variants[var_id]
            if not variant_obj.get('compounds'):
                continue

            updated_compounds = self.update_variant_compounds(variant_obj, variants)
            variant_obj['compounds'] = updated_compounds

        LOG.debug("Compounds updated")

        return variants

    def update_mongo_compound_variants(self, bulk):
        """Update the compound information for a bulk of variants in the database

            Args:
                bulk(dict): {'_id': scout.models.Variant}

        """
        requests = []
        for var_id in bulk:
            var_obj = bulk[var_id]
            if not var_obj.get('compounds'):
                continue
            # Add a request to update compounds
            operation = pymongo.UpdateOne(
                {'_id': var_obj['_id']},
                {
                    '$set': {
                        'compounds': var_obj['compounds']
                    }
                })
            requests.append(operation)

        if not requests:
            return

        try:
            self.variant_collection.bulk_write(requests, ordered=False)
        except BulkWriteError as err:
            LOG.warning("Updating compounds failed")
            raise err

    def update_case_compounds(self, case_obj, build='37'):
        """Update the compounds for a case

        Loop over all coding intervals to get coordinates for all potential compound positions.
        Update all variants within a gene with a bulk operation.
        """

        case_id = case_obj['_id']
        # Possible categories 'snv', 'sv', 'str', 'cancer':
        categories = set()
        # Possible variant types 'clinical', 'research':
        variant_types = set()

        for file_type in FILE_TYPE_MAP:
            if case_obj.get('vcf_files',{}).get(file_type):
                categories.add(FILE_TYPE_MAP[file_type]['category'])
                variant_types.add(FILE_TYPE_MAP[file_type]['variant_type'])

        coding_intervals = self.get_coding_intervals(build=build)
        # Loop over all intervals
        for chrom in CHROMOSOMES:
            intervals = coding_intervals.get(chrom, IntervalTree())
            for var_type in variant_types:
                for category in categories:
                    LOG.info("Updating compounds on chromosome:{0}, type:{1}, category:{2} for case:{3}".format(
                             chrom, var_type, category, case_id))

                    # Fetch all variants from a chromosome
                    query  = {
                        'variant_type': var_type,
                        'chrom': chrom,
                    }

                    # Get all variants from the database of the specific type
                    variant_objs = self.variants(
                        case_id=case_id,
                        query=query,
                        category=category,
                        nr_of_variants=-1,
                        sort_key='position'
                    )

                    # Initiate a bulk
                    bulk = {}
                    current_region = None
                    special = False

                    # Loop over the variants and check if they are in a coding region
                    for var_obj in variant_objs:
                        var_id = var_obj['_id']
                        var_chrom = var_obj['chromosome']
                        var_start = var_obj['position']
                        var_end = var_obj['end'] + 1

                        update_bulk = True
                        new_region = None

                        # Check if the variant is in a coding region
                        genomic_regions = coding_intervals.get(var_chrom, IntervalTree()).search(var_start, var_end)

                        # If the variant is in a coding region
                        if genomic_regions:
                            # We know there is data here so get the interval id
                            new_region = genomic_regions.pop().data

                        if new_region and (new_region == current_region):
                            # If the variant is in the same region as previous
                            # we add it to the same bulk
                            update_bulk = False

                        current_region = new_region

                        # If the variant is not in a current region we update the compounds
                        # from the previous region, if any. Otherwise continue
                        if update_bulk and bulk:
                            self.update_compounds(bulk)
                            self.update_mongo_compound_variants(bulk)
                            bulk = {}

                        if new_region:
                            bulk[var_id] = var_obj

                    if not bulk:
                        continue

                    self.update_compounds(bulk)
                    self.update_mongo_compound_variants(bulk)

        LOG.info("All compounds updated")
        return

    def load_variant(self, variant_obj):
        """Load a variant object

        Args:
            variant_obj(dict)

        Returns:
            inserted_id
        """
        # LOG.debug("Loading variant %s", variant_obj['_id'])
        try:
            result = self.variant_collection.insert_one(variant_obj)
        except DuplicateKeyError as err:
            raise IntegrityError("Variant %s already exists in database", variant_obj['_id'])
        return result

    def upsert_variant(self, variant_obj):
        """Load a variant object, if the object already exists update compounds.

        Args:
            variant_obj(dict)

        Returns:
            result
        """
        LOG.debug("Upserting variant %s", variant_obj['_id'])
        try:
            result = self.variant_collection.insert_one(variant_obj)
        except DuplicateKeyError as err:
            LOG.debug("Variant %s already exists in database", variant_obj['_id'])
            result = self.variant_collection.find_one_and_update(
                {'_id': variant_obj['_id']},
                {
                    '$set': {
                        'compounds': variant_obj.get('compounds',[])
                    }
                }
            )
            variant = self.variant_collection.find_one({'_id': variant_obj['_id']})
        return result

    def load_variant_bulk(self, variants):
        """Load a bulk of variants

        Args:
            variants(iterable(scout.models.Variant))

        Returns:
            object_ids
        """
        if not len(variants) > 0:
            return

        LOG.debug("Loading variant bulk")
        try:
            result = self.variant_collection.insert_many(variants)
        except (DuplicateKeyError, BulkWriteError) as err:
            # If the bulk write is wrong there are probably some variants already existing
            # In the database. So insert each variant
            for var_obj in variants:
                try:
                    self.upsert_variant(var_obj)
                except IntegrityError as err:
                    pass

        return

    def _load_variants(self, variants, variant_type, case_obj, individual_positions, rank_threshold,
                       institute_id, build=None, rank_results_header=None, vep_header=None,
                       category='snv', sample_info = None):
        """Perform the loading of variants

        This is the function that loops over the variants, parse them and build the variant
        objects so they are ready to be inserted into the database.

        """
        build = build or '37'
        genes = [gene_obj for gene_obj in self.all_genes(build=build)]
        gene_to_panels = self.gene_to_panels(case_obj)
        hgncid_to_gene = self.hgncid_to_gene(genes=genes)
        genomic_intervals = self.get_coding_intervals(genes=genes)

        LOG.info("Start inserting {0} {1} variants into database".format(variant_type, category))
        start_insertion = datetime.now()
        start_five_thousand = datetime.now()
        # These are the number of parsed varaints
        nr_variants = 0
        # These are the number of variants that meet the criteria and gets inserted
        nr_inserted = 0
        # This is to keep track of blocks of inserted variants
        inserted = 1

        nr_bulks = 0

        # We want to load batches of variants to reduce the number of network round trips
        bulk = {}
        current_region = None

        for nr_variants, variant in enumerate(variants):
            # All MT variants are loaded
            mt_variant = 'MT' in variant.CHROM
            rank_score = parse_rank_score(variant.INFO.get('RankScore'), case_obj['_id'])

            # Check if the variant should be loaded at all
            # if rank score is None means there are no rank scores annotated, all variants will be loaded
            # Otherwise we load all variants above a rank score treshold
            # Except for MT variants where we load all variants
            if (rank_score is None) or (rank_score > rank_threshold) or mt_variant:
                nr_inserted += 1
                # Parse the vcf variant
                parsed_variant = parse_variant(
                    variant=variant,
                    case=case_obj,
                    variant_type=variant_type,
                    rank_results_header=rank_results_header,
                    vep_header=vep_header,
                    individual_positions=individual_positions,
                    category=category,
                )

                # Build the variant object
                variant_obj = build_variant(
                    variant=parsed_variant,
                    institute_id=institute_id,
                    gene_to_panels=gene_to_panels,
                    hgncid_to_gene=hgncid_to_gene,
                    sample_info=sample_info
                )

                # Check if the variant is in a genomic region
                var_chrom = variant_obj['chromosome']
                var_start = variant_obj['position']
                # We need to make sure that the interval has a length > 0
                var_end = variant_obj['end'] + 1
                var_id = variant_obj['_id']
                # If the bulk should be loaded or not
                load = True
                new_region = None

                genomic_regions = genomic_intervals.get(var_chrom, IntervalTree()).search(var_start, var_end)

                # If the variant is in a coding region
                if genomic_regions:
                    # We know there is data here so get the interval id
                    new_region = genomic_regions.pop().data
                    # If the variant is in the same region as previous
                    # we add it to the same bulk
                    if new_region == current_region:
                        load = False

                # This is the case where the variant is intergenic
                else:
                    # If the previous variant was also intergenic we add the variant to the bulk
                    if not current_region:
                        load = False
                    # We need to have a max size of the bulk
                    if len(bulk) > 10000:
                        load = True
                # Load the variant object
                if load:
                    # If the variant bulk contains coding variants we want to update the compounds
                    if current_region:
                        self.update_compounds(bulk)
                    try:
                        # Load the variants
                        self.load_variant_bulk(list(bulk.values()))
                        nr_bulks += 1
                    except IntegrityError as error:
                        pass
                    bulk = {}

                current_region = new_region
                bulk[var_id] = variant_obj

                if (nr_variants != 0 and nr_variants % 5000 == 0):
                    LOG.info("%s variants parsed", str(nr_variants))
                    LOG.info("Time to parse variants: %s",
                                (datetime.now() - start_five_thousand))
                    start_five_thousand = datetime.now()

                if (nr_inserted != 0 and (nr_inserted * inserted) % (1000 * inserted) == 0):
                    LOG.info("%s variants inserted", nr_inserted)
                    inserted += 1
        # If the variants are in a coding region we update the compounds
        if current_region:
            self.update_compounds(bulk)

        # Load the final variant bulk
        self.load_variant_bulk(list(bulk.values()))
        nr_bulks += 1
        LOG.info("All variants inserted, time to insert variants: {0}".format(
            datetime.now() - start_insertion))

        if nr_variants:
            nr_variants += 1
        LOG.info("Nr variants parsed: %s", nr_variants)
        LOG.info("Nr variants inserted: %s", nr_inserted)
        LOG.debug("Nr bulks inserted: %s", nr_bulks)

        return nr_inserted

    def load_variants(self, case_obj, variant_type='clinical', category='snv',
                      rank_threshold=None, chrom=None, start=None, end=None,
                      gene_obj=None, build='37'):
        """Load variants for a case into scout.

        Load the variants for a specific analysis type and category into scout.
        If no region is specified, load all variants above rank score threshold
        If region or gene is specified, load all variants from that region
        disregarding variant rank(if not specified)

        Args:
            case_obj(dict): A case from the scout database
            variant_type(str): 'clinical' or 'research'. Default: 'clinical'
            category(str): 'snv', 'str' or 'sv'. Default: 'snv'
            rank_threshold(float): Only load variants above this score. Default: 0
            chrom(str): Load variants from a certain chromosome
            start(int): Specify the start position
            end(int): Specify the end position
            gene_obj(dict): A gene object from the database

        Returns:
            nr_inserted(int)
        """
        # We need the institute object
        institute_id = self.institute(institute_id=case_obj['owner'])['_id']
        nr_inserted = 0

        variant_file = None
        if variant_type == 'clinical':
            if category == 'snv':
                variant_file = case_obj['vcf_files'].get('vcf_snv')
            elif category == 'sv':
                variant_file = case_obj['vcf_files'].get('vcf_sv')
            elif category == 'str':
                LOG.debug('Attempt to load STR VCF.')
                variant_file = case_obj['vcf_files'].get('vcf_str')
            elif category == 'cancer':
                # Currently this implies a paired tumor normal
                variant_file = case_obj['vcf_files'].get('vcf_cancer')
        elif variant_type == 'research':
            if category == 'snv':
                variant_file = case_obj['vcf_files'].get('vcf_snv_research')
            elif category == 'sv':
                variant_file = case_obj['vcf_files'].get('vcf_sv_research')
            elif category == 'cancer':
                variant_file = case_obj['vcf_files'].get('vcf_cancer_research')

        if not variant_file:
            raise SyntaxError("Vcf file does not seem to exist")

        # Check if there are any variants in file
        try:
            vcf_obj = VCF(variant_file)
            var = next(vcf_obj)
        except StopIteration as err:
            LOG.warning("Variant file %s does not include any variants", variant_file)
            return nr_inserted
        # We need to reload the file
        vcf_obj = VCF(variant_file)

        # Parse the neccessary headers from vcf file
        rank_results_header = parse_rank_results_header(vcf_obj)
        vep_header = parse_vep_header(vcf_obj)

        # This is a dictionary to tell where ind are in vcf
        individual_positions = {}
        for i, ind in enumerate(vcf_obj.samples):
            individual_positions[ind] = i

        # Dictionary for cancer analysis
        sample_info = {}
        if category == 'cancer':
            for ind in case_obj['individuals']:
                if ind['phenotype'] == 2:
                    sample_info[ind['individual_id']] = 'case'
                else:
                    sample_info[ind['individual_id']] = 'control'

        # Check if a region scould be uploaded
        region = ""
        if gene_obj:
            chrom = gene_obj['chromosome']
            # Add same padding as VEP
            start = max(gene_obj['start'] - 5000, 0)
            end = gene_obj['end'] + 5000
        if chrom:
            # We want to load all variants in the region regardless of rank score
            rank_threshold = rank_threshold or -1000
            if not (start and end):
                raise SyntaxError("Specify chrom start and end")
            region = "{0}:{1}-{2}".format(chrom, start, end)
        else:
            rank_threshold = rank_threshold or 0

        variants = vcf_obj(region)

        try:
            nr_inserted = self._load_variants(
                variants=variants,
                variant_type=variant_type,
                case_obj=case_obj,
                individual_positions=individual_positions,
                rank_threshold=rank_threshold,
                institute_id=institute_id,
                build=build,
                rank_results_header=rank_results_header,
                vep_header=vep_header,
                category=category,
                sample_info = sample_info
            )
        except Exception as error:
            LOG.exception('unexpected error')
            LOG.warning("Deleting inserted variants")
            self.delete_variants(case_obj['_id'], variant_type)
            raise error

        self.update_variant_rank(case_obj, variant_type, category=category)

        return nr_inserted


import logging
from datetime import datetime

from bson import ObjectId
import pymongo

from scout.constants import CASE_STATUSES, REV_ACMG_MAP

LOG = logging.getLogger(__name__)

class CaseEventHandler(object):
    """Class to handle case events for the mongo adapter"""

    def assign(self, institute, case, user, link):
        """Assign a user to a case.

        This function will create an Event to log that a person has been assigned
        to a case. Also the user will be added to case "assignees".

        Arguments:
            institute (dict): A institute
            case (dict): A case
            user (dict): A User object
            link (str): The url to be used in the event

        Returns:
            updated_case(dict)
        """
        LOG.info("Creating event for assigning {0} to {1}"
                    .format(user['name'].encode('utf-8'), case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='assign',
            subject=case['display_name']
        )
        LOG.info("Updating {0} to be assigned with {1}"
                    .format(case['display_name'], user['name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$addToSet': {'assignees': user['_id']}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        return updated_case

    def unassign(self, institute, case, user, link):
        """Unassign a user from a case.

        This function will create an Event to log that a person has been
        unassigned from a case. Also the user will be removed from case
        "assignees".

        Arguments:
            institute (dict): A Institute object
            case (dict): A Case object
            user (dict): A User object (Should this be a user id?)
            link (dict): The url to be used in the event

        Returns:
            updated_case (dict): The updated case
        """
        LOG.info("Creating event for unassigning {0} from {1}".format(
            user['name'], case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='unassign',
            subject=case['display_name']
        )

        LOG.info("Updating {0} to be unassigned with {1}".format(
            case['display_name'], user['name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$pull': {'assignees': user['_id']}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def update_status(self, institute, case, user, status, link):
        """Update the status of a case.

        This function will create an Event to log that a user have updated the
        status of a case. Also the status of the case will be updated.
        Status could be anyone of:
            ("prioritized", "inactive", "active", "solved", "archived")

        Arguments:
            institute (dict): A Institute object
            case (dict): A Case object
            user (dict): A User object
            status (str): The new status of the case
            link (str): The url to be used in the event

        Returns:
            updated_case
        """

        if status not in CASE_STATUSES:
            LOG.warning("Status {0} is invalid".format(status))
            return None

        LOG.info("Creating event for updating status of {0} to {1}".format(
            case['display_name'], status))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='status',
            subject=case['display_name']
        )

        LOG.info("Updating {0} to status {1}".format(case['display_name'], status))
        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$set': {'status': status}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def update_synopsis(self, institute, case, user, link, content=""):
        """Create an Event for updating the synopsis for a case.

            This function will create an Event and update the synopsis for
             a case.

        Arguments:
            institute (dict): A Institute object
            case (dict): A Case object
            user (dict): A User object
            link (str): The url to be used in the event
            content (str): The content for what should be added to the synopsis

        Returns:
            updated_case
        """
        LOG.info("Creating event for updating the synopsis for case" \
                    " {0}".format(case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='synopsis',
            subject=case['display_name'],
            content=content
        )

        LOG.info("Updating the synopsis for case {0}".format(
            case['display_name']))
        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$set': {'synopsis': content}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def archive_case(self, institute, case, user, link):
        """Create an event for archiving a case.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event

        Returns:
            updated_case (dict)
        """
        LOG.info("Creating event for archiving case {0}".format(
            case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='archive',
            subject=case['display_name'],
        )

        LOG.info("Change status for case {0} to 'archived'".format(
            case['display_name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$set': {'status': 'archived'}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def open_research(self, institute, case, user, link):
        """Create an event for opening the research list a case.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event

        Returns:
            updated_case(dict)
        """
        LOG.info("Creating event for opening research for case"
                    " {0}".format(case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='open_research',
            subject=case['display_name'],
        )

        LOG.info("Set research_requested for case {0} to True".format(
            case['display_name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$set': {'research_requested': True}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def request_rerun(self, institute, case, user, link):
        """Request a case to be re-analyzed.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event

        Return:
            updated_case
        """
        if case.get('rerun_requested'):
            raise ValueError('rerun already pending')

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='rerun',
            subject=case['display_name']
        )

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$set': {'rerun_requested': True}
            },
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def share(self, institute, case, collaborator_id, user, link):
        """Share a case with a new institute.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            collaborator_id (str): A instute id
            user (dict): A User object
            link (str): The url to be used in the event

        Return:
            updated_case
        """
        if collaborator_id in case.get('collaborators', []):
            raise ValueError('new customer is already a collaborator')

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='share',
            subject=collaborator_id
        )

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$push': {'collaborators': collaborator_id}
            },
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def unshare(self, institute, case, collaborator_id, user, link):
        """Revoke access for a collaborator for a case.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            collaborator_id (str): A instute id
            user (dict): A User object
            link (str): The url to be used in the event

        Return:
            updated_case

        """
        if collaborator_id not in case['collaborators']:
            raise ValueError("collaborator doesn't have access to case")

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='unshare',
            subject=collaborator_id
        )

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$pull': {'collaborators': collaborator_id}
            },
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def diagnose(self, institute, case, user, link, level, omim_id, remove=False):
        """Diagnose a case using OMIM ids.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            level (str): choices=('phenotype','gene')

        Return:
            updated_case

        """
        if level == 'phenotype':
            case_key = 'diagnosis_phenotypes'
        elif level == 'gene':
            case_key = 'diagnosis_genes'
        else:
            raise TypeError('wrong level')

        diagnosis_list = case.get(case_key, [])
        omim_number = int(omim_id.split(':')[-1])

        updated_case = None
        if remove and omim_number in diagnosis_list:
            updated_case = self.case_collection.find_one_and_update(
                {'_id': case['_id']},
                {'$pull': {case_key: omim_number}},
                return_document=pymongo.ReturnDocument.AFTER
            )

        elif omim_number not in diagnosis_list:
            updated_case = self.case_collection.find_one_and_update(
                {'_id': case['_id']},
                {'$push': {case_key: omim_number}},
                return_document=pymongo.ReturnDocument.AFTER
            )

        if updated_case:
            self.create_event(
                institute=institute,
                case=case,
                user=user,
                link=link,
                category='case',
                verb='update_diagnosis',
                subject=case['display_name'],
                content=omim_id
            )

        return updated_case

    def add_cohort(self, institute, case, user, link, tag):
        """Add a cohort tag to case

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            tag (str): The cohort tag to be added

        Return:
            updated_case(dict)
        """
        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='add_cohort',
            subject=link
        )

        LOG.info("Adding cohort tag {0} to {1}"
                    .format(tag, case['display_name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$addToSet': {'cohorts': tag},
            },
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def remove_cohort(self, institute, case, user, link, tag):
        """Remove a cohort tag from case

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            tag (str): The cohort tag to be removed

        Return:
            updated_case(dict)
        """
        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='remove_cohort',
            subject=case['display_name'],
        )

        LOG.info("Removing cohort tag {0} to {1}"
                    .format(tag, case['display_name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$pull': {'cohorts': tag},
            },
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def mark_checked(self, institute, case, user, link,
                     unmark=False):
        """Mark a case as checked from an analysis point of view.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            unmark (bool): If case should ve unmarked

        Return:
            updated_case
        """

        LOG.info("Updating checked status of {}"
                    .format(case['display_name']))

        status = 'not checked' if unmark else 'checked'
        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='check_case',
            subject=status
        )

        LOG.info("Updating {0}'s checked status {1}"
                    .format(case['display_name'], status))
        analysis_checked = False if unmark else True
        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$set': {'analysis_checked': analysis_checked}
            },
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def update_default_panels(self, institute_obj, case_obj, user_obj, link, panel_objs):
        """Update default panels for a case.

        Arguments:
            institute_obj (dict): A Institute object
            case_obj (dict): Case object
            user_obj (dict): A User object
            link (str): The url to be used in the event
            panel_objs (list(dict)): List of panel objs

        Return:
            updated_case(dict)

        """
        self.create_event(
            institute=institute_obj,
            case=case_obj,
            user=user_obj,
            link=link,
            category='case',
            verb='update_default_panels',
            subject=case_obj['display_name'],
        )

        LOG.info("Update default panels for {}".format(case_obj['display_name']))

        panel_ids = [panel['_id'] for panel in panel_objs]

        for existing_panel in case_obj['panels']:
            if existing_panel['panel_id'] in panel_ids:
                existing_panel['is_default'] = True
            else:
                existing_panel['is_default'] = False

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case_obj['_id']},
            {
                '$set': {'panels': case_obj['panels']}
            },
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")

        return updated_case


import logging
from datetime import datetime

from bson import ObjectId
import pymongo

from scout.constants import REV_ACMG_MAP

SANGER_OPTIONS = ['True positive', 'False positive', 'Not validated']

LOG = logging.getLogger(__name__)

class VariantEventHandler(object):
    """Class to handle variant events for the mongo adapter"""

    def pin_variant(self, institute, case, user, link, variant):
        """Create an event for pinning a variant.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object

        Returns:
            updated_case
        """
        LOG.info("Creating event for pinning variant {0}".format(
            variant['display_name']))

        # add variant to list of pinned references in the case model
        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$push': {'suspects': variant['_id']}},
            return_document=pymongo.ReturnDocument.AFTER
        )

        kwargs = dict(
            institute=institute,
            case=case,
            user=user,
            link=link,
            verb='pin',
            variant=variant,
            subject=variant['display_name'],
        )
        self.create_event(category='variant', **kwargs)
        self.create_event(category='case', **kwargs)

        return updated_case

    def unpin_variant(self, institute, case, user, link, variant):
        """Create an event for unpinning a variant.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object

        Returns:
            updated_case(dict)
        """
        LOG.info("Creating event for unpinning variant {0}".format(
            variant['display_name']))

        LOG.info("Remove variant from list of references in the case" \
                    " model")

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$pull': {'suspects': variant['_id']}},
            return_document=pymongo.ReturnDocument.AFTER
        )

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='unpin',
            variant=variant,
            subject=variant['display_name'],
        )

        return updated_case

    def order_verification(self, institute, case, user, link, variant):
        """Create an event for a variant verification for a variant
        and an event for a variant verification for a case

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object

        Returns:
            updated_variant(dict)
        """
        LOG.info("Creating event for ordering validation for variant" \
                    " {0}".format(variant['display_name']))

        updated_variant = self.variant_collection.find_one_and_update(
            {'_id': variant['_id']},
            {'$set': {'sanger_ordered': True}},
            return_document=pymongo.ReturnDocument.AFTER
        )

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='sanger',
            variant=variant,
            subject=variant['display_name'],
        )

        LOG.info("Creating event for ordering sanger for case" \
                    " {0}".format(case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='sanger',
            variant=variant,
            subject=variant['display_name'],
        )
        return updated_variant

    def cancel_verification(self, institute, case, user, link, variant):
        """Create an event for cancellation of a verification of a variant

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object

        Returns:
            updated_variant(dict)
        """
        LOG.info("Creating event for cancellation of ordering sanger for variant" \
                    " {0}".format(variant['display_name']))

        updated_variant = self.variant_collection.find_one_and_update(
            {'_id': variant['_id']},
            {'$set': {'sanger_ordered': False}},
            return_document=pymongo.ReturnDocument.AFTER
        )

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='cancel_sanger',
            variant=variant,
            subject=variant['display_name'],
        )

        LOG.info("Creating event for cancellation of ordering verification for case" \
                    " {0}".format(case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='cancel_sanger',
            variant=variant,
            subject=variant['display_name'],
        )
        return updated_variant

    def sanger_ordered(self, institute_id=None, user_id=None):
        """Get all variants with validations ever ordered.

        Args:
            institute_id(str) : The id of an institute
            user_id(str) : The id of an user

        Returns:
            sanger_ordered(list) : a list of dictionaries, each with "case_id" as keys and list of variant ids as values
        """
        query = {'$match': {
                '$and': [
                    {'verb': 'sanger'},
                ],
            }}

        if institute_id:
            query['$match']['$and'].append({'institute': institute_id})
        if user_id:
            query['$match']['$and'].append({'user_id': user_id})

        # Get all sanger ordered variants grouped by case_id
        results = self.event_collection.aggregate([
            query,
            {'$group': {
                '_id': "$case",
                'vars': {'$addToSet' : '$variant_id'}
            }}
        ])

        sanger_ordered =  [item for item in results]
        return sanger_ordered

    def validate(self, institute, case, user, link, variant, validate_type):
        """Mark validation status for a variant.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object
            validate_type(str): The outcome of validation.
                                choices=('True positive', 'False positive')

        Returns:
            updated_variant(dict)
        """
        if not validate_type in SANGER_OPTIONS:
            LOG.warning("Invalid validation string: %s", validate_type)
            LOG.info("Validation options: %s", ', '.join(SANGER_OPTIONS))
            return

        updated_variant = self.variant_collection.find_one_and_update(
            {'_id': variant['_id']},
            {'$set': {'validation': validate_type}},
            return_document=pymongo.ReturnDocument.AFTER
        )

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='validate',
            variant=variant,
            subject=variant['display_name'],
        )
        return updated_variant

    def mark_causative(self, institute, case, user, link, variant):
        """Create an event for marking a variant causative.

        Arguments:
          institute (dict): A Institute object
          case (dict): Case object
          user (dict): A User object
          link (str): The url to be used in the event
          variant (variant): A variant object

        Returns:
            updated_case(dict)
        """
        display_name = variant['display_name']
        LOG.info("Mark variant {0} as causative in the case {1}".format(
            display_name, case['display_name']))

        LOG.info("Adding variant to causatives in case {0}".format(
            case['display_name']))

        LOG.info("Marking case {0} as solved".format(
            case['display_name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$push': {'causatives': variant['_id']},
                '$set': {'status': 'solved'}
            },
            return_document=pymongo.ReturnDocument.AFTER
        )

        LOG.info("Creating case event for marking {0}" \
                    " causative".format(variant['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='mark_causative',
            variant=variant,
            subject=variant['display_name'],
        )

        LOG.info("Creating variant event for marking {0}" \
                    " causative".format(case['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='mark_causative',
            variant=variant,
            subject=variant['display_name'],
        )
        return updated_case

    def unmark_causative(self, institute, case, user, link, variant):
        """Create an event for unmarking a variant causative

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object

        Returns:
            updated_case(dict)

        """
        display_name = variant['display_name']
        LOG.info("Remove variant {0} as causative in case {1}".format(
            display_name, case['display_name']))

        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {
                '$pull': {'causatives': variant['_id']},
            },
            return_document=pymongo.ReturnDocument.AFTER
        )

        # mark the case as active again
        if len(updated_case.get('causatives', [])) == 0:
            LOG.info("Marking case as 'active'")
            updated_case = self.case_collection.find_one_and_update(
                {'_id': case['_id']},
                {
                    '$set': {'status': 'active'}
                },
                return_document=pymongo.ReturnDocument.AFTER
            )

        LOG.info("Creating events for unmarking variant {0} " \
                    "causative".format(display_name))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='case',
            verb='unmark_causative',
            variant=variant,
            subject=variant['display_name'],
        )

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='unmark_causative',
            variant=variant,
            subject=variant['display_name'],
        )

        return updated_case

    def update_manual_rank(self, institute, case, user, link, variant,
                           manual_rank):
        """Create an event for updating the manual rank of a variant

          This function will create a event and update the manual rank
          of the variant.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object
            manual_rank (int): The new manual rank

        Return:
            updated_variant

        """
        LOG.info("Creating event for updating manual rank for "
                    "variant {0}".format(variant['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='manual_rank',
            variant=variant,
            subject=variant['display_name'],
        )

        if manual_rank:
            LOG.info("Setting manual rank to {0} for variant {1}"
                        .format(manual_rank, variant['display_name']))
            action = '$set'
        else:
            LOG.info("Reset manual rank from {0} for variant {1}"
                        .format(variant.get('manual_rank', 'NA'), variant['display_name']))
            action = '$unset'

        updated_variant = self.variant_collection.find_one_and_update(
            {'_id': variant['_id']},
            {action: {'manual_rank': manual_rank}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Variant updated")
        return updated_variant

    def update_dismiss_variant(self, institute, case, user, link, variant,
                               dismiss_variant):
        """Create an event for updating the manual dismiss variant entry

          This function will create a event and update the dismiss variant
          field of the variant.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object
            dismiss_variant (list): The new dismiss variant list

        Return:
            updated_variant

        """
        LOG.info("Creating event for updating dismiss variant for "
                    "variant {0}".format(variant['display_name']))

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='dismiss_variant',
            variant=variant,
            subject=variant['display_name'],
        )

        if dismiss_variant:
            LOG.info("Setting dismiss variant to {0} for variant {1}"
                        .format(dismiss_variant, variant['display_name']))
            action = '$set'
        else:
            LOG.info("Reset dismiss variant from {0} for variant {1}"
                        .format(variant['dismiss_variant'], variant['display_name']))
            action = '$unset'

        updated_variant = self.variant_collection.find_one_and_update(
            {'_id': variant['_id']},
            {action: {'dismiss_variant': dismiss_variant}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Variant updated")
        return updated_variant

    def update_mosaic_tags(self, institute, case, user, link, variant,
                               mosaic_tags):
        """Create an event for updating the mosaicism tags for a variant

          This function will create a event and update the mosaicism tags of the variant.

        Arguments:
            institute (dict): A Institute object
            case (dict): Case object
            user (dict): A User object
            link (str): The url to be used in the event
            variant (dict): A variant object
            mosaic_tags (list): The new mosaic tags list

        Return:
            updated_variant

        """
        LOG.info("Creating event for updating mosaic tags for variant %s", variant['display_name'])

        self.create_event(
            institute=institute,
            case=case,
            user=user,
            link=link,
            category='variant',
            verb='mosaic_tags',
            variant=variant,
            subject=variant['display_name'],
        )

        if mosaic_tags:
            LOG.info("Setting mosaic tags to {0} for variant {1}"
                        .format(mosaic_tags, variant['display_name']))
            action = '$set'
        else:
            LOG.info("Reset dismiss variant from {0} for variant {1}"
                        .format(variant['mosaic_tags'], variant['display_name']))
            action = '$unset'

        updated_variant = self.variant_collection.find_one_and_update(
            {'_id': variant['_id']},
            {action: {'mosaic_tags': mosaic_tags}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Variant updated")
        return updated_variant

    def update_acmg(self, institute_obj, case_obj, user_obj, link, variant_obj, acmg_str):
        """Create an event for updating the ACMG classification of a variant.

        Arguments:
            institute_obj (dict): A Institute object
            case_obj (dict): Case object
            user_obj (dict): A User object
            link (str): The url to be used in the event
            variant_obj (dict): A variant object
            acmg_str (str): The new ACMG classification string

        Returns:
            updated_variant
        """
        self.create_event(
            institute=institute_obj,
            case=case_obj,
            user=user_obj,
            link=link,
            category='variant',
            verb='acmg',
            variant=variant_obj,
            subject=variant_obj['display_name'],
        )
        LOG.info("Setting ACMG to {} for: {}".format(acmg_str, variant_obj['display_name']))

        if acmg_str is None:
            updated_variant = self.variant_collection.find_one_and_update(
                {'_id': variant_obj['_id']},
                {'$unset': {'acmg_classification': 1}},
                return_document=pymongo.ReturnDocument.AFTER
            )
        else:
            updated_variant = self.variant_collection.find_one_and_update(
                {'_id': variant_obj['_id']},
                {'$set': {'acmg_classification': REV_ACMG_MAP[acmg_str]}},
                return_document=pymongo.ReturnDocument.AFTER
            )

        LOG.debug("Variant updated")
        return updated_variant


from scout.utils.md5 import generate_md5_key

def parse_ids(chrom, pos, ref, alt, case_id, variant_type):
    """Construct the necessary ids for a variant

    Args:
        chrom(str): Variant chromosome
        pos(int): Variant position
        ref(str): Variant reference
        alt(str): Variant alternative
        case_id(str): Unique case id
        variant_type(str): 'clinical' or 'research'

    Returns:
        ids(dict): Dictionary with the relevant ids
    """
    ids = {}
    pos = str(pos)

    ids['simple_id'] = parse_simple_id(chrom, pos, ref, alt)
    ids['variant_id'] = parse_variant_id(chrom, pos, ref, alt, variant_type)
    ids['display_name'] = parse_display_name(chrom, pos, ref, alt, variant_type)
    ids['document_id'] = parse_document_id(chrom, pos, ref, alt, variant_type, case_id)

    return ids

def parse_simple_id(chrom, pos, ref, alt):
    """Parse the simple id for a variant

    Simple id is used as a human readable reference for a position, it is
    in no way unique.

    Args:
        chrom(str)
        pos(str)
        ref(str)
        alt(str)

    Returns:
        simple_id(str): The simple human readable variant id
    """
    return '_'.join([chrom, pos, ref, alt])

def parse_variant_id(chrom, pos, ref, alt, variant_type):
    """Parse the variant id for a variant

    variant_id is used to identify variants within a certain type of
    analysis. It is not human readable since it is a md5 key.

    Args:
        chrom(str)
        pos(str)
        ref(str)
        alt(str)
        variant_type(str): 'clinical' or 'research'

    Returns:
        variant_id(str): The variant id converted to md5 string
    """
    return generate_md5_key([chrom, pos, ref, alt, variant_type])

def parse_display_name(chrom, pos, ref, alt, variant_type):
    """Parse the variant id for a variant

    This is used to display the variant in scout.

    Args:
        chrom(str)
        pos(str)
        ref(str)
        alt(str)
        variant_type(str): 'clinical' or 'research'

    Returns:
        variant_id(str): The variant id in human readable format
    """
    return '_'.join([chrom, pos, ref, alt, variant_type])

def parse_document_id(chrom, pos, ref, alt, variant_type, case_id):
    """Parse the unique document id for a variant.

    This will always be unique in the database.

    Args:
        chrom(str)
        pos(str)
        ref(str)
        alt(str)
        variant_type(str): 'clinical' or 'research'
        case_id(str): unqiue family id

    Returns:
        document_id(str): The unique document id in an md5 string
    """
    return generate_md5_key([chrom, pos, ref, alt, variant_type, case_id])


# -*- coding: utf-8 -*-
import logging

import click

from scout.utils.handle import get_file_handle
from scout.parse.panel import parse_genes

log = logging.getLogger(__name__)

@click.command('convert', short_help='Convert gene panels')
@click.argument('panel', type=click.File('r'))
@click.pass_context
def convert(context, panel):
    """Convert a gene panel with hgnc symbols to a new one with hgnc ids."""
    adapter = context.obj['adapter']
    new_header = ["hgnc_id","hgnc_symbol","disease_associated_transcripts",
                 "reduced_penetrance", "genetic_disease_models", "mosaicism",
                 "database_entry_version"]

    genes = parse_genes(panel)

    adapter.add_hgnc_id(genes)

    click.echo("#{0}".format('\t'.join(new_header)))
    for gene in genes:
        if gene.get('hgnc_id'):
            print_info = []
            for head in new_header:
                print_info.append(str(gene[head]) if gene.get(head) else '')

            click.echo('\t'.join(print_info))


# -*- coding: utf-8 -*-
from copy import deepcopy
import logging
import datetime

from pprint import pprint as pp

import pymongo

from scout.parse.case import parse_case
from scout.build.case import build_case
from scout.parse.variant.ids import parse_document_id

from scout.exceptions import IntegrityError, ConfigError

LOG = logging.getLogger(__name__)

class CaseHandler(object):
    """Part of the pymongo adapter that handles cases and institutes"""

    def cases(self, owner=None, collaborator=None, query=None, skip_assigned=False,
              has_causatives=False, reruns=False, finished=False,
              research_requested=False, is_research=False, status=None,
              phenotype_terms=False, pinned=False, cohort=False, name_query=None,
              yield_query=False):
        """Fetches all cases from the backend.

        Args:
            collaborator(str): If collaborator should be considered
            owner(str): Query cases for specified case owner only
            query(dict): If a specific query is used
            skip_assigned(bool)
            has_causatives(bool)
            reruns(bool)
            finished(bool)
            research_requested(bool)
            is_research(bool)
            status(str)
            phenotype_terms(bool): Fetch all cases with phenotype terms
            pinned(bool): Fetch all cases with pinned variants
            name_query(str): Could be hpo term, HPO-group, user, part of display name,
                             part of inds or part of synopsis
            yield_query(bool): If true, only return mongo query dict for use in
                                compound querying.

        Returns:
            Cases ordered by date.
            If yield_query is True, does not pose query to db;
                instead returns corresponding query dict
                that can be reused in compound queries or for testing.
        """
        LOG.debug("Fetch all cases")
        query = query or {}

        # Prioritize when both owner and collaborator params are present
        if collaborator and owner:
            collaborator = None

        if collaborator:
            LOG.debug("Use collaborator {0}".format(collaborator))
            query['collaborators'] = collaborator

        if owner:
            LOG.debug("Use owner {0}".format(owner))
            query['owner'] = owner

        if skip_assigned:
            query['assignees'] = {'$exists': False}

        if has_causatives:
            query['causatives'] = {'$exists': True, '$ne': []}

        if reruns:
            query['rerun_requested'] = True

        if status:
            query['status'] = status

        elif finished:
            query['status'] = {'$in': ['solved', 'archived']}

        if research_requested:
            query['research_requested'] = True

        if is_research:
            query['is_research'] = {'$exists': True, '$eq': True}

        if phenotype_terms:
            query['phenotype_terms'] = {'$exists': True, '$ne': []}

        if pinned:
            query['suspects'] = {'$exists': True, '$ne': []}

        if cohort:
            query['cohorts'] = {'$exists': True, '$ne': []}

        if name_query:
            name_value = name_query.split(':')[-1] # capture ant value provided after query descriptor
            users = self.user_collection.find({'name': {'$regex': name_query, '$options': 'i'}})
            if users.count() > 0:
                query['assignees'] = {'$in': [user['email'] for user in users]}
            elif name_query.startswith('HP:'):
                LOG.debug("HPO case query")
                if name_value:
                    query['phenotype_terms.phenotype_id'] = name_query
                else: # query for cases with no HPO terms
                    query['$or'] = [ {'phenotype_terms' : {'$size' : 0}}, {'phenotype_terms' : {'$exists' : False}} ]
            elif name_query.startswith('PG:'):
                LOG.debug("PG case query")
                if name_value:
                    phenotype_group_query = name_query.replace('PG:', 'HP:')
                    query['phenotype_groups.phenotype_id'] = phenotype_group_query
                else: # query for cases with no phenotype groups
                    query['$or'] = [ {'phenotype_groups' : {'$size' : 0}}, {'phenotype_groups' : {'$exists' : False}} ]
            elif name_query.startswith('synopsis:'):
                if name_value:
                    query['$text']={'$search':name_value}
                else: # query for cases with missing synopsis
                    query['synopsis'] = ''
            elif name_query.startswith('cohort:'):
                query['cohorts'] = name_value
            elif name_query.startswith('panel:'):
                query['panels'] = {'$elemMatch': {'panel_name': name_value,
                                    'is_default': True }}
            elif name_query.startswith('status:'):
                status_query = name_query.replace('status:','')
                query['status'] = status_query
            elif name_query.startswith('is_research'):
                query['is_research'] = {'$exists': True, '$eq': True}
            else:
                query['$or'] = [
                    {'display_name': {'$regex': name_query}},
                    {'individuals.display_name': {'$regex': name_query}},
                ]

        if yield_query:
            return query

        LOG.info("Get cases with query {0}".format(query))
        return self.case_collection.find(query).sort('updated_at', -1)

    def nr_cases(self, institute_id=None):
        """Return the number of cases

        This function will change when we migrate to 3.7.1

        Args:
            collaborator(str): Institute id

        Returns:
            nr_cases(int)
        """
        query = {}

        if institute_id:
            query['collaborators'] = institute_id

        LOG.debug("Fetch all cases with query {0}".format(query))
        nr_cases = self.case_collection.find(query).count()

        return nr_cases

    def update_dynamic_gene_list(self, case, hgnc_symbols=None, hgnc_ids=None,
                                 phenotype_ids=None, build='37'):
        """Update the dynamic gene list for a case

        Adds a list of dictionaries to case['dynamic_gene_list'] that looks like

        {
            hgnc_symbol: str,
            hgnc_id: int,
            description: str
        }

        Arguments:
            case (dict): The case that should be updated
            hgnc_symbols (iterable): A list of hgnc_symbols
            hgnc_ids (iterable): A list of hgnc_ids

        Returns:
            updated_case(dict)
        """
        dynamic_gene_list = []
        res = []
        if hgnc_ids:
            LOG.info("Fetching genes by hgnc id")
            res = self.hgnc_collection.find({'hgnc_id': {'$in': hgnc_ids}, 'build': build})
        elif hgnc_symbols:
            LOG.info("Fetching genes by hgnc symbols")
            res = []
            for symbol in hgnc_symbols:
                for gene_obj in self.gene_by_alias(symbol=symbol, build=build):
                    res.append(gene_obj)

        for gene_obj in res:
            dynamic_gene_list.append(
                {
                    'hgnc_symbol': gene_obj['hgnc_symbol'],
                    'hgnc_id': gene_obj['hgnc_id'],
                    'description': gene_obj['description'],
                }
            )

        LOG.info("Update dynamic gene panel for: %s", case['display_name'])
        updated_case = self.case_collection.find_one_and_update(
            {'_id': case['_id']},
            {'$set': {'dynamic_gene_list': dynamic_gene_list,
                      'dynamic_panel_phenotypes': phenotype_ids or []}},
            return_document=pymongo.ReturnDocument.AFTER
        )
        LOG.debug("Case updated")
        return updated_case

    def case(self, case_id=None, institute_id=None, display_name=None):
        """Fetches a single case from database

        Use either the _id or combination of institute_id and display_name

        Args:
            case_id(str): _id for a caes
            institute_id(str):
            display_name(str)

        Yields:
            A single Case
        """
        query = {}
        if case_id:
            query['_id'] = case_id
            LOG.info("Fetching case %s", case_id)
        else:
            if not (institute_id and display_name):
                raise ValueError("Have to provide both institute_id and display_name")
            LOG.info("Fetching case %s institute %s", display_name, institute_id)
            query['owner'] = institute_id
            query['display_name'] = display_name

        return self.case_collection.find_one(query)

    def case_ind(self, ind_id):
        """Fetch cases based on an individual id.

        Args:
            ind_id(str)

        Returns:
            cases(pymongo.cursor): The cases with a matching ind_id
        """

        return self.case_collection.find({'individuals.disply_name': ind_id})

    def delete_case(self, case_id=None, institute_id=None, display_name=None):
        """Delete a single case from database

        Args:
            institute_id(str)
            case_id(str)

        Returns:
            case_obj(dict): The case that was deleted
        """
        query = {}
        if case_id:
            query['_id'] = case_id
            LOG.info("Deleting case %s", case_id)
        else:
            if not (institute_id and display_name):
                raise ValueError("Have to provide both institute_id and display_name")
            LOG.info("Deleting case %s institute %s", display_name, institute_id)
            query['owner'] = institute_id
            query['display_name'] = display_name

        result = self.case_collection.delete_one(query)
        return result

    def load_case(self, config_data, update=False):
        """Load a case into the database

        Check if the owner and the institute exists.

        Args:
            config_data(dict): A dictionary with all the necessary information
            update(bool): If existing case should be updated

        Returns:
            case_obj(dict)
        """
        # Check that the owner exists in the database
        institute_obj = self.institute(config_data['owner'])
        if not institute_obj:
            raise IntegrityError("Institute '%s' does not exist in database" % config_data['owner'])

        # Parse the case information
        parsed_case = parse_case(config=config_data)
        # Build the case object
        case_obj = build_case(parsed_case, self)
        # Check if case exists with old case id
        old_caseid = '-'.join([case_obj['owner'], case_obj['display_name']])
        old_case = self.case(old_caseid)
        if old_case:
            LOG.info("Update case id for existing case: %s -> %s", old_caseid, case_obj['_id'])
            self.update_caseid(old_case, case_obj['_id'])
            update = True

        # Check if case exists in database
        existing_case = self.case(case_obj['_id'])
        if existing_case and not update:
            raise IntegrityError("Case %s already exists in database" % case_obj['_id'])

        files = [
            {'file_name': 'vcf_snv', 'variant_type': 'clinical', 'category': 'snv'},
            {'file_name': 'vcf_sv', 'variant_type': 'clinical', 'category': 'sv'},
            {'file_name': 'vcf_cancer', 'variant_type': 'clinical', 'category': 'cancer'},
            {'file_name': 'vcf_str', 'variant_type': 'clinical', 'category': 'str'}
        ]

        try:
            for vcf_file in files:
                # Check if file exists
                if not case_obj['vcf_files'].get(vcf_file['file_name']):
                    LOG.debug("didn't find {}, skipping".format(vcf_file['file_name']))
                    continue

                variant_type = vcf_file['variant_type']
                category = vcf_file['category']
                if update:
                    self.delete_variants(
                        case_id=case_obj['_id'],
                        variant_type=variant_type,
                        category=category
                    )
                self.load_variants(
                    case_obj=case_obj,
                    variant_type=variant_type,
                    category=category,
                    rank_threshold=case_obj.get('rank_score_threshold', 0),
                )

        except (IntegrityError, ValueError, ConfigError, KeyError) as error:
            LOG.warning(error)

        if existing_case and update:
            self.update_case(case_obj)
        else:
            LOG.info('Loading case %s into database', case_obj['display_name'])
            self._add_case(case_obj)

        return case_obj

    def _add_case(self, case_obj):
        """Add a case to the database
           If the case already exists exception is raised

            Args:
                case_obj(Case)
        """
        if self.case(case_obj['_id']):
            raise IntegrityError("Case %s already exists in database" % case_obj['_id'])

        return self.case_collection.insert_one(case_obj)

    def update_case(self, case_obj):
        """Update a case in the database

        The following will be updated:
            - collaborators: If new collaborators these will be added to the old ones
            - analysis_date: Is updated to the new date
            - analyses: The new analysis date will be added to old runs
            - individuals: There could be new individuals
            - updated_at: When the case was updated in the database
            - rerun_requested: Is set to False since that is probably what happened
            - panels: The new gene panels are added
            - genome_build: If there is a new genome build
            - genome_version: - || -
            - rank_model_version: If there is a new rank model
            - madeline_info: If there is a new pedigree
            - vcf_files: paths to the new files
            - has_svvariants: If there are new svvariants
            - has_strvariants: If there are new strvariants
            - multiqc: If there's an updated multiqc report location
            - mme_submission: If case was submitted to MatchMaker Exchange

            Args:
                case_obj(dict): The new case information

            Returns:
                updated_case(dict): The updated case information
        """
        # Todo: rename to match the intended purpose

        LOG.info("Updating case {0}".format(case_obj['_id']))
        old_case = self.case_collection.find_one(
            {'_id': case_obj['_id']}
        )
        updated_case = self.case_collection.find_one_and_update(
            {'_id': case_obj['_id']},
            {
                '$addToSet': {
                    'collaborators': {'$each': case_obj['collaborators']},
                    'analyses': {
                        'date': old_case['analysis_date'],
                        'delivery_report': old_case.get('delivery_report')
                    }
                },
                '$set': {
                    'analysis_date': case_obj['analysis_date'],
                    'delivery_report': case_obj.get('delivery_report'),
                    'individuals': case_obj['individuals'],
                    'updated_at': datetime.datetime.now(),
                    'rerun_requested': False,
                    'panels': case_obj.get('panels', []),
                    'genome_build': case_obj.get('genome_build', '37'),
                    'genome_version': case_obj.get('genome_version'),
                    'rank_model_version': case_obj.get('rank_model_version'),
                    'madeline_info': case_obj.get('madeline_info'),
                    'vcf_files': case_obj.get('vcf_files'),
                    'has_svvariants': case_obj.get('has_svvariants'),
                    'has_strvariants': case_obj.get('has_strvariants'),
                    'is_research': case_obj.get('is_research', False),
                    'research_requested': case_obj.get('research_requested', False),
                    'multiqc': case_obj.get('multiqc'),
                    'mme_submission': case_obj.get('mme_submission'),
                }
            },
            return_document=pymongo.ReturnDocument.AFTER
        )

        LOG.info("Case updated")
        return updated_case

    def replace_case(self, case_obj):
        """Replace a existing case with a new one

        Keeps the object id

        Args:
            case_obj(dict)

        Returns:
            updated_case(dict)
        """
        # Todo: Figure out and describe when this method destroys a case if invoked instead of
        # update_case
        LOG.info("Saving case %s", case_obj['_id'])
        # update updated_at of case to "today"

        case_obj['updated_at'] = datetime.datetime.now(),

        updated_case = self.case_collection.find_one_and_replace(
            {'_id': case_obj['_id']},
            case_obj,
            return_document=pymongo.ReturnDocument.AFTER
        )

        return updated_case

    def update_caseid(self, case_obj, family_id):
        """Update case id for a case across the database.

        This function is used when a case is a rerun or updated for another reason.

        Args:
            case_obj(dict)
            family_id(str): The new family id

        Returns:
            new_case(dict): The updated case object

        """
        new_case = deepcopy(case_obj)
        new_case['_id'] = family_id

        # update suspects and causatives
        for case_variants in ['suspects', 'causatives']:
            new_variantids = []
            for variant_id in case_obj.get(case_variants, []):
                case_variant = self.variant(variant_id)
                if not case_variant:
                    continue
                new_variantid = get_variantid(case_variant, family_id)
                new_variantids.append(new_variantid)
            new_case[case_variants] = new_variantids

        # update ACMG
        for acmg_obj in self.acmg_collection.find({'case_id': case_obj['_id']}):
            LOG.info("update ACMG classification: %s", acmg_obj['classification'])
            acmg_variant = self.variant(acmg_obj['variant_specific'])
            new_specific_id = get_variantid(acmg_variant, family_id)
            self.acmg_collection.find_one_and_update(
                {'_id': acmg_obj['_id']},
                {'$set': {'case_id': family_id, 'variant_specific': new_specific_id}},
            )

        # update events
        institute_obj = self.institute(case_obj['owner'])
        for event_obj in self.events(institute_obj, case=case_obj):
            LOG.info("update event: %s", event_obj['verb'])
            self.event_collection.find_one_and_update(
                {'_id': event_obj['_id']},
                {'$set': {'case': family_id}},
            )

        # insert the updated case
        self.case_collection.insert_one(new_case)
        # delete the old case
        self.case_collection.find_one_and_delete({'_id': case_obj['_id']})
        return new_case

def get_variantid(variant_obj, family_id):
    """Create a new variant id.

    Args:
        variant_obj(dict)
        family_id(str)

    Returns:
        new_id(str): The new variant id
    """
    new_id = parse_document_id(
        chrom=variant_obj['chromosome'],
        pos=str(variant_obj['position']),
        ref=variant_obj['reference'],
        alt=variant_obj['alternative'],
        variant_type=variant_obj['variant_type'],
        case_id=family_id,
    )
    return new_id


# -*- coding: utf-8 -*-
import logging

from bson import ObjectId
import pymongo

from scout.utils.acmg import get_acmg
from scout.build.acmg import build_evaluation

log = logging.getLogger(__name__)

class ACMGHandler(object):

    def submit_evaluation(self, variant_obj, user_obj, institute_obj, case_obj, link, criteria):
        """Submit an evaluation to the database

        Get all the relevant information, build a evaluation_obj

        Args:
            variant_obj(dict)
            user_obj(dict)
            institute_obj(dict)
            case_obj(dict)
            link(str): variant url
            criteria(list(dict)):

                [
            {
            'term': str,
            'comment': str,
            'links': list(str)
            },
            .
            .
        ]
        """
        variant_specific = variant_obj['_id']
        variant_id = variant_obj['variant_id']
        user_id = user_obj['_id']
        user_name = user_obj.get('name', user_obj['_id'])
        institute_id = institute_obj['_id']
        case_id = case_obj['_id']

        evaluation_terms = [evluation_info['term'] for evluation_info in criteria]

        classification = get_acmg(evaluation_terms)

        evaluation_obj = build_evaluation(
            variant_specific=variant_specific,
            variant_id=variant_id,
            user_id=user_id,
            user_name=user_name,
            institute_id=institute_id,
            case_id=case_id,
            classification=classification,
            criteria=criteria
        )

        self._load_evaluation(evaluation_obj)

        # Update the acmg classification for the variant:
        self.update_acmg(institute_obj, case_obj, user_obj, link, variant_obj, classification)
        return classification

    def _load_evaluation(self, evaluation_obj):
        """Load a evaluation object into the database"""
        res = self.acmg_collection.insert_one(evaluation_obj)
        return res

    def delete_evaluation(self, evaluation_obj):
        """Delete an evaluation from the database

        Args:
            evaluation_obj(dict)

        """
        self.acmg_collection.delete_one({'_id': evaluation_obj['_id']})

    def get_evaluation(self, evaluation_id):
        """Get a single evaluation from the database

        Args:
            evaluation_id(str)

        """
        return self.acmg_collection.find_one({'_id': ObjectId(evaluation_id)})

    def get_evaluations(self, variant_obj):
        """Return all evaluations for a certain variant.

        Args:
            variant_obj (dict): variant dict from the database

        Returns:
            pymongo.cursor: database cursor
        """
        query = dict(variant_id=variant_obj['variant_id'])
        res = self.acmg_collection.find(query).sort([('created_at', pymongo.DESCENDING)])
        return res


import logging

from pprint import pprint as pp

from pandas import DataFrame

LOG = logging.getLogger(__name__)

def parse_transcripts(transcript_lines):
    """Parse and massage the transcript information

    There could be multiple lines with information about the same transcript.
    This is why it is necessary to parse the transcripts first and then return a dictionary
    where all information has been merged.

    Args:
        transcript_lines(): This could be an iterable with strings or a pandas.DataFrame

    Returns:
        parsed_transcripts(dict): Map from enstid -> transcript info
    """
    LOG.info("Parsing transcripts")
    # Parse the transcripts, we need to check if it is a request or a file handle
    if isinstance(transcript_lines, DataFrame):
        transcripts = parse_ensembl_transcript_request(transcript_lines)
    else:
        transcripts = parse_ensembl_transcripts(transcript_lines)

    # Since there can be multiple lines with information about the same transcript
    # we store transcript information in a dictionary for now
    parsed_transcripts = {}
    # Loop over the parsed transcripts
    for tx in transcripts:
        tx_id = tx['ensembl_transcript_id']
        ens_gene_id = tx['ensembl_gene_id']

        # Check if the transcript has been added
        # If not, create a new transcript
        if not tx_id in parsed_transcripts:
            tx_info = {
                'chrom': tx['chrom'],
                'transcript_start': tx['transcript_start'],
                'transcript_end': tx['transcript_end'],
                'mrna': set(),
                'mrna_predicted': set(),
                'nc_rna': set(),
                'ensembl_gene_id': ens_gene_id,
                'ensembl_transcript_id': tx_id,
            }
            parsed_transcripts[tx_id] = tx_info

        tx_info = parsed_transcripts[tx_id]
        # Add the ref seq information
        if tx.get('refseq_mrna_predicted'):
            tx_info['mrna_predicted'].add(tx['refseq_mrna_predicted'])
        if tx.get('refseq_mrna'):
            tx_info['mrna'].add(tx['refseq_mrna'])
        if tx.get('refseq_ncrna'):
            tx_info['nc_rna'].add(tx['refseq_ncrna'])

    return parsed_transcripts

def parse_ensembl_gene_request(result):
    """Parse a dataframe with ensembl gene information

    Args:
        res(pandas.DataFrame)

    Yields:
        gene_info(dict)
    """
    LOG.info("Parsing genes from request")

    for index, row in result.iterrows():
        # print(index, row)
        ensembl_info = {}

        # Pandas represents missing data with nan which is a float
        if type(row['hgnc_symbol']) is float:
            # Skip genes without hgnc information
            continue

        ensembl_info['chrom'] = row['chromosome_name']
        ensembl_info['gene_start'] = int(row['start_position'])
        ensembl_info['gene_end'] = int(row['end_position'])
        ensembl_info['ensembl_gene_id'] = row['ensembl_gene_id']
        ensembl_info['hgnc_symbol'] = row['hgnc_symbol']

        hgnc_id = row['hgnc_id']

        if type(hgnc_id) is float:
            hgnc_id = int(hgnc_id)
        else:
            hgnc_id = int(hgnc_id.split(':')[-1])

        ensembl_info['hgnc_id'] = hgnc_id

        yield ensembl_info

def parse_ensembl_transcript_request(result):
    """Parse a dataframe with ensembl transcript information

    Args:
        res(pandas.DataFrame)

    Yields:
        transcript_info(dict)
    """
    LOG.info("Parsing transcripts from request")

    keys = [
        'chrom',
        'ensembl_gene_id',
        'ensembl_transcript_id',
        'transcript_start',
        'transcript_end',
        'refseq_mrna',
        'refseq_mrna_predicted',
        'refseq_ncrna',
    ]
    # for res in result.itertuples():
    for index, row in result.iterrows():
        ensembl_info = {}

        ensembl_info['chrom'] = str(row['chromosome_name'])
        ensembl_info['ensembl_gene_id'] = row['ensembl_gene_id']
        ensembl_info['ensembl_transcript_id'] = row['ensembl_transcript_id']

        ensembl_info['transcript_start'] = int(row['transcript_start'])
        ensembl_info['transcript_end'] = int(row['transcript_end'])

        # Check if refseq data is annotated
        # Pandas represent missing data with nan
        for key in keys[-3:]:
            if type(row[key]) is float:
                ensembl_info[key] = None
            else:
                ensembl_info[key] = row[key]
        yield ensembl_info

def parse_ensembl_line(line, header):
    """Parse an ensembl formated line

        Args:
            line(list): A list with ensembl gene info
            header(list): A list with the header info

        Returns:
            ensembl_info(dict): A dictionary with the relevant info
    """
    line = line.rstrip().split('\t')
    header = [head.lower() for head in header]
    raw_info = dict(zip(header, line))

    ensembl_info = {}

    for word in raw_info:
        value = raw_info[word]

        if not value:
            continue

        if 'chromosome' in word:
            ensembl_info['chrom'] = value

        if 'gene' in word:
            if 'id' in word:
                ensembl_info['ensembl_gene_id'] = value
            elif 'start' in word:
                ensembl_info['gene_start'] = int(value)
            elif 'end' in word:
                ensembl_info['gene_end'] = int(value)

        if 'hgnc symbol' in word:
            ensembl_info['hgnc_symbol'] = value
        if "gene name" in word:
            ensembl_info['hgnc_symbol'] = value

        if 'hgnc id' in word:
            ensembl_info['hgnc_id'] = int(value.split(':')[-1])

        if 'transcript' in word:
            if 'id' in word:
                ensembl_info['ensembl_transcript_id'] = value
            elif 'start' in word:
                ensembl_info['transcript_start'] = int(value)
            elif 'end' in word:
                ensembl_info['transcript_end'] = int(value)

        if 'exon' in word:
            if 'start' in word:
                ensembl_info['exon_start'] = int(value)
            elif 'end' in word:
                ensembl_info['exon_end'] = int(value)
            elif 'rank' in word:
                ensembl_info['exon_rank'] = int(value)

        if 'utr' in word:

            if 'start' in word:
                if '5' in word:
                    ensembl_info['utr_5_start'] = int(value)
                elif '3' in word:
                    ensembl_info['utr_3_start'] = int(value)
            elif 'end' in word:
                if '5' in word:
                    ensembl_info['utr_5_end'] = int(value)
                elif '3' in word:
                    ensembl_info['utr_3_end'] = int(value)

        if 'strand' in word:
            ensembl_info['strand'] = int(value)

        if 'refseq' in word:
            if 'mrna' in word:
                if 'predicted' in word:
                    ensembl_info['refseq_mrna_predicted'] = value
                else:
                    ensembl_info['refseq_mrna'] = value

            if 'ncrna' in word:
                ensembl_info['refseq_ncrna'] = value

    return ensembl_info

def parse_ensembl_genes(lines):
    """Parse lines with ensembl formated genes

        This is designed to take a biomart dump with genes from ensembl.
        Mandatory columns are:
        'Gene ID' 'Chromosome' 'Gene Start' 'Gene End' 'HGNC symbol

        Args:
            lines(iterable(str)): An iterable with ensembl formated genes
        Yields:
            ensembl_gene(dict): A dictionary with the relevant information
    """
    LOG.info("Parsing ensembl genes from file")
    header = []
    for index, line in enumerate(lines):

        # File allways start with a header line
        if index == 0:
            header = line.rstrip().split('\t')
            continue
        # After that each line represents a gene

        yield parse_ensembl_line(line, header)

def parse_ensembl_transcripts(lines):
    """Parse lines with ensembl formated transcripts

        This is designed to take a biomart dump with transcripts from ensembl.
        Mandatory columns are:
        'Gene ID' 'Transcript ID' 'Transcript Start' 'Transcript End' 'RefSeq mRNA'

        Args:
            lines(iterable(str)): An iterable with ensembl formated genes
        Yields:
            ensembl_gene(dict): A dictionary with the relevant information
    """
    header = []
    LOG.info("Parsing ensembl genes from file")
    for index, line in enumerate(lines):

        # File allways start with a header line
        if index == 0:
            header = line.rstrip().split('\t')
        # After that each line represents a transcript
        else:
            yield parse_ensembl_line(line, header)

def parse_ensembl_exons(lines):
    """Parse lines with ensembl formated exons

        This is designed to take a biomart dump with exons from ensembl.
        Check documentation for spec for download

        Args:
            lines(iterable(str)): An iterable with ensembl formated exons
        Yields:
            ensembl_gene(dict): A dictionary with the relevant information
    """
    header = []
    LOG.debug("Parsing ensembl exons...")
    for index, line in enumerate(lines):

        # File allways start with a header line
        if index == 0:
            header = line.rstrip().split('\t')
            continue

        exon_info = parse_ensembl_line(line, header)
        chrom = exon_info['chrom']
        start = exon_info['exon_start']
        end = exon_info['exon_end']
        transcript = exon_info['ensembl_transcript_id']
        gene = exon_info['ensembl_gene_id']

        rank = exon_info['exon_rank']
        strand = exon_info['strand']

        # Recalculate start and stop (taking UTR regions into account for end exons)
        if strand == 1:
            # highest position: start of exon or end of 5' UTR
            # If no 5' UTR make sure exon_start is allways choosen
            start = max(start, exon_info.get('utr_5_end') or -1)
            # lowest position: end of exon or start of 3' UTR
            end = min(end, exon_info.get('utr_3_start') or float('inf'))
        elif strand == -1:
            # highest position: start of exon or end of 3' UTR
            start = max(start, exon_info.get('utr_3_end') or -1)
            # lowest position: end of exon or start of 5' UTR
            end = min(end, exon_info.get('utr_5_start') or float('inf'))

        exon_id = "-".join([chrom, str(start), str(end)])

        if start > end:
            raise ValueError("ERROR: %s" % exon_id)
        data = {
            "exon_id": exon_id,
            "chrom": chrom,
            "start": start,
            "end": end,
            "transcript": transcript,
            "gene": gene,
            "rank": rank,
        }

        yield data

def parse_ensembl_exon_request(result):
    """Parse a dataframe with ensembl exon information

    Args:
        res(pandas.DataFrame)

    Yields:
        gene_info(dict)
    """
    keys = [
        'chrom',
        'gene',
        'transcript',
        'exon_id',
        'exon_chrom_start',
        'exon_chrom_end',
        '5_utr_start',
        '5_utr_end',
        '3_utr_start',
        '3_utr_end',
        'strand',
        'rank'
    ]

    # for res in result.itertuples():
    for res in zip(result['Chromosome/scaffold name'],
                   result['Gene stable ID'],
                   result['Transcript stable ID'],
                   result['Exon stable ID'],
                   result['Exon region start (bp)'],
                   result['Exon region end (bp)'],
                   result["5' UTR start"],
                   result["5' UTR end"],
                   result["3' UTR start"],
                   result["3' UTR end"],
                   result["Strand"],
                   result["Exon rank in transcript"]):
        ensembl_info = dict(zip(keys, res))

        # Recalculate start and stop (taking UTR regions into account for end exons)
        if ensembl_info['strand'] == 1:
            # highest position: start of exon or end of 5' UTR
            # If no 5' UTR make sure exon_start is allways choosen
            start = max(ensembl_info['exon_chrom_start'], ensembl_info['5_utr_end'] or -1)
            # lowest position: end of exon or start of 3' UTR
            end = min(ensembl_info['exon_chrom_end'], ensembl_info['3_utr_start'] or float('inf'))
        elif ensembl_info['strand'] == -1:
            # highest position: start of exon or end of 3' UTR
            start = max(ensembl_info['exon_chrom_start'], ensembl_info['3_utr_end'] or -1)
            # lowest position: end of exon or start of 5' UTR
            end = min(ensembl_info['exon_chrom_end'], ensembl_info['5_utr_start'] or float('inf'))

        ensembl_info['start'] = start
        ensembl_info['end'] = end

        yield ensembl_info

        # if type(ensembl_info['hgnc_symbol']) is float:
        #     # Skip genes without hgnc information
        #     continue
        # ensembl_info['gene_start'] = int(ensembl_info['gene_start'])
        # ensembl_info['gene_end'] = int(ensembl_info['gene_end'])
        #
        # if type(ensembl_info['hgnc_id']) is float:
        #     ensembl_info['hgnc_id'] = int(ensembl_info['hgnc_id'])
        # else:
        #     ensembl_info['hgnc_id'] = int(ensembl_info['hgnc_id'].split(':')[-1])
        #
        # yield ensembl_info