# #################################################
# SNP INPUT MIXIN
# #################################################


class SNPInputMixin:
    """
    Provides methods for managing SNP filters in the database, allowing for
    union and intersection operations on the SNP dataset to either initialize
    or refine the dataset by matching with a provided list of SNP identifiers.

    Methods:
    - `unionInputSNPs`: Adds new SNPs from a provided list to the SNP dataset
        in the specified database, initializing the dataset if it hasn't been
        set up already.
    - `intersectInputSNPs`: Refines the SNP dataset by intersecting it with a
        provided list, keeping only matching SNPs.

    If the SNP filter is not yet initialized, this method will first
    initialize it with a union.
    """

    def unionInputSNPs(self, db, snps, errorCallback=None):
        """
        Adds a list of SNPs to a specific SNP filter within an `snp` table in
        a specified database, using `rs` identifiers and labels, and handles
        cases of ambiguous or merged SNPs.

        Parameters:
        - db: Name of the database where SNPs will be added. Example: `user`.
        - snps: List of tuples `[(rs, extra), ...]`, where each `rs` is the
            SNP identifier and `extra` holds additional information.
        - errorCallback: Optional function called in case of an error during
            SNP addition.

        Operation:
        - Logs the start of the SNP addition process to the SNP filter in the
            specified database (`db`).
        - Calls `prepareTableForUpdate` to drop indexes from the `snp` table
            before insertion to improve performance.
        - Defines and executes an `INSERT INTO` SQL query to add SNPs
            generated by `generateMergedFilteredSNPs`.
        - Updates the `tally` dictionary to track how many SNPs were added and
            how many identifiers were merged or ambiguous.
        - Logs the number of SNPs added, showing counts of merged (`merge`)
            and ambiguous (`many`) SNPs, if any.
        - Increments the `snp` filter counter in the `db` database to mark the
            addition.

        Returns:
        - None. The method adds SNPs to the specified table and logs
            information about the process.

        This method enables adding and managing SNPs in an `snp` table,
        applying a merging process and handling ambiguous SNPs, providing
        flexibility to work with complex SNP data.
        """
        # snps=[ (rs,extra), ... ]
        self.logPush("adding to %s SNP filter ...\n" % db)
        cursor = self._loki._db.cursor()

        self.prepareTableForUpdate(db, "snp")
        sql = (
            "INSERT INTO `%s`.`snp` (label,extra,rs) VALUES ('rs'||?1,?2,?3)"
            % db  # noqa E501
        )  # noqa E501
        tally = dict()
        cursor.executemany(
            sql, self.generateMergedFilteredSNPs(snps, tally, errorCallback)
        )

        if tally.get("many"):
            self.logPop(
                "... OK: added %d SNPs (%d RS#s merged, %d ambiguous)\n"
                % (
                    tally["match"] + tally["merge"] - tally["many"],
                    tally["merge"],
                    tally["many"],
                )
            )
        else:
            self.logPop(
                "... OK: added %d SNPs (%d RS#s merged)\n"
                % (tally["match"] + tally["merge"], tally["merge"])
            )
        self._inputFilters[db]["snp"] += 1

    def intersectInputSNPs(self, db, snps, errorCallback=None):
        """
        Reduces an SNP filter in a specific database table, retaining only
        SNPs present in the provided list. If the SNP filter is not yet
        initialized, performs a union rather than an intersection.

        Parameters:
        - `db`: The name of the database where the `snp` table is located.
        - `snps`: A list of tuples `[(rs, extra), ...]`, where `rs` is the SNP
            identifier, and `extra` contains additional information.
        - `errorCallback`: Optional function called in case of an error during
            SNP processing.

        Operation:
        - If the SNP filter has not yet been initialized for the database
            (`_inputFilters[db]['snp']` is 0), it calls `unionInputSNPs` to
            initialize the filter.
        - Otherwise:
            - Logs the start of the reduction process.
            - Calls `prepareTableForQuery` to ensure the `snp` table is ready
                for querying.
            - Sets all SNPs in the table to "not retained" (`flag = 0`).
            - Counts the number of SNPs before reduction (`numBefore`).
            - Updates the `snp` table, setting `flag = 1` for SNPs matching
                the `rs` identifiers in the `snps` list.
            - Deletes SNPs that are not in the provided list, i.e., with
                `flag = 0`.
            - Counts the number of SNPs removed (`numDrop`).
            - Logs the final count of retained SNPs, discarded SNPs, and
                merged identifiers (`merge`).
            - Increments the SNP filter counter for the database.

        Returns:
        - Nothing. The method performs an intersection of SNPs in the
            specified table and logs the outcome.

        This function manages an SNP set in a database table by performing an
        intersection between existing SNPs and a provided list, retaining only
        matches.
        """
        # snps=[ (rs,extra), ... ]
        if not self._inputFilters[db]["snp"]:
            return self.unionInputSNPs(db, snps, errorCallback)
        self.logPush("reducing %s SNP filter ...\n" % db)
        cursor = self._loki._db.cursor()

        self.prepareTableForQuery(db, "snp")
        cursor.execute("UPDATE `%s`.`snp` SET flag = 0" % db)
        numBefore = cursor.getconnection().changes()
        sql = (
            "UPDATE `%s`.`snp` SET flag = 1 WHERE (1 OR ?1 OR ?2) AND rs = ?3"
            % db  # noqa E501
        )  # noqa E501
        tally = dict()
        # we don't have to do ambiguous snp filtering here because we're only
        # reducing what's already loaded
        cursor.executemany(
            sql, self._loki.generateCurrentRSesByRSes(snps, tally)
        )  # noqa E501
        cursor.execute("DELETE FROM `%s`.`snp` WHERE flag = 0" % db)
        numDrop = cursor.getconnection().changes()

        self.logPop(
            "... OK: kept %d SNPs (%d dropped, %d RS#s merged)\n"
            % (numBefore - numDrop, numDrop, tally["merge"])
        )
        self._inputFilters[db]["snp"] += 1
