"""
Support for masking potential CpG sites in *pairwise* alignments.
"""

from bx.align.sitemask import Masker
from ._cpg import (
    list_cpg,
    list_cpg_restricted,
    list_non_cpg,
)


# Restricted.  Only mask out sites that are defitely CpG
class Restricted(Masker):
    def __init__(self, mask="?"):
        self.mask = mask
        self.masked = 0
        self.total = 0

    def __call__(self, block):
        if not block:
            return block
        if len(block.components) < 2:
            return
        cpglist = list_cpg_restricted(block.components[0].text.upper(), block.components[1].text.upper())

        # now we have a fast list of CpG columns, iterate/mask
        self.masked += len(cpglist)
        self.total += len(block.components[0].text)
        for component in block.components:
            component.text = mask_columns(cpglist, component.text, self.mask)

        return block


# Inclusive. Mask out all sites that are not non-CpG sites.


class Inclusive(Masker):
    def __init__(self, mask="?"):
        self.mask = mask
        self.masked = 0
        self.total = 0

    def __call__(self, block):
        if not block:
            return block
        if len(block.components) < 2:
            return
        cpglist = list_cpg(block.components[0].text.upper(), block.components[1].text.upper())

        self.masked += len(cpglist)
        self.total += len(block.components[0].text)
        for component in block.components:
            component.text = mask_columns(cpglist, component.text, self.mask)

        return block


# Mak nonCpG sites


class nonCpG(Masker):
    def __init__(self, mask="?"):
        self.mask = mask
        self.masked = 0
        self.total = 0

    def __call__(self, block):
        if not block:
            return block
        if len(block.components) < 2:
            return
        noncpglist = list_non_cpg(block.components[0].text.upper(), block.components[1].text.upper())

        # now we have a fast list of non-CpG columns, iterate/mask
        self.masked += len(noncpglist)
        self.total += len(block.components[0].text)
        for component in block.components:
            component.text = mask_columns(noncpglist, component.text, self.mask)

        return block


def mask_columns(masklist, text, mask):
    templist = []
    for position in masklist:
        if text[position] != "-":
            templist.append(position)
    templist.append(len(text))  # Add the end of the text
    # cut string
    newtext = []
    c = 0
    for position in templist:
        newtext.append(text[c:position])
        c = position + 1  # Gaps have len = 1
    joinedtext = mask.join(newtext)
    return joinedtext