import numpy as np
import pandas as pd


class InvalidPositionException(Exception):
    def __init__(self, message):
        super().__init__(message)


class PositionEncoder:
    __slots__ = ("chromosomes", "chromosome_sizes", "sizes", "offsets")

    CHR_NOT_SUPPORTED = -1
    ERROR_WRONG_CHR_POSITION = -2

    def __init__(self):
        self.chromosomes = [
            "1",
            "2",
            "3",
            "4",
            "5",
            "6",
            "7",
            "8",
            "9",
            "10",
            "11",
            "12",
            "13",
            "14",
            "15",
            "16",
            "17",
            "18",
            "19",
            "20",
            "21",
            "22",
            "X",
            "Y",
            "M",
        ]

        self.chromosome_sizes = {
            "1": 248956422,
            "2": 242193529,
            "3": 198295559,
            "4": 190214555,
            "5": 181538259,
            "6": 170805979,
            "7": 159345973,
            "8": 145138636,
            "9": 138394717,
            "10": 133797422,
            "11": 135086622,
            "12": 133275309,
            "13": 114364328,
            "14": 107043718,
            "15": 101991189,
            "16": 90338345,
            "17": 83257441,
            "18": 80373285,
            "19": 58617616,
            "20": 64444167,
            "21": 46709983,
            "22": 50818468,
            "X": 156040895,
            "Y": 57227415,
            "M": 16569,
        }
        self.sizes = [
            self.chromosome_sizes.get(chr_name, 0) for chr_name in self.chromosomes
        ]
        self.offsets = [sum(self.sizes[:i]) for i in range(len(self.chromosomes))]

    def encode(self, chromosome: str, position: int):
        return self.encode_1_based_without_throw(chromosome, position)

    def encode_1_based_without_throw(self, chromosome, position):
        chr_name = self._chr(chromosome)
        try:
            chr_index = self.chromosomes.index(chr_name)
        except ValueError:
            # chromosome not supported
            return self.CHR_NOT_SUPPORTED
        else:
            size = self.sizes[chr_index]
            if position > size or position < 0:
                print(f"Wrong position for {chromosome}: {position}")
                return self.ERROR_WRONG_CHR_POSITION
            offset = self.offsets[chr_index]
            return position - 1 + offset

    @staticmethod
    def _chr(chromosome):
        # This method corresponds to ChromosomeNameNormalizer.chr(chromosome) in the Java code
        return chromosome  # You may need to implement your normalization logic here


# Singleton instance of PositionEncoder
_position_encoder_instance = PositionEncoder()


def encode_vcf_position_gbid(chromosome, position):
    return _position_encoder_instance.encode_1_based_without_throw(chromosome, position)


class VariantIdEncoder:
    __slots__ = ("positionEncoder",)

    RADIX = 36
    MAX_DEL_LENGTH = (2**7) - 1
    MAX_INS_LENGTH = 9

    MASK_HASH = 0b1000000000000000000000000000000000000000000000000000000000000000
    MASK_NULL = 0b0100000000000000000000000000000000000000000000000000000000000000
    MASK_POS = 0b0011111111111111111111111111111111000000000000000000000000000000
    MASK_C_HASH = 0b0000000000000000000000000000000000100000000000000000000000000000
    MASK_INS_LEN = 0b0000000000000000000000000000000000011110000000000000000000000000
    MASK_DEL = 0b0000000000000000000000000000000000000001111111000000000000000000
    MASK_INS = 0b0000000000000000000000000000000000000000000000111111111111111111
    MASK_CHANGE_HASH_VALUE = (
        0b0000000000000000000000000000000000011111111111111111111111111111
    )

    # Pre-computed values for optimization
    BASE_ENCODING = {"A": 0, "C": 1, "G": 2, "T": 3}
    VALID_BASES = frozenset("ACGNT")

    # Pre-computed shift amounts
    SHIFT_POS = 30
    SHIFT_INS_LEN = 25
    SHIFT_DEL = 18
    SHIFT_C_HASH = 29

    def __init__(self):
        self.positionEncoder = PositionEncoder()

    def hash_function(self, input_string: str, seed: int = 104729) -> int:
        try:
            # try with a fast c-implementation ...
            import mmh3 as mmh3
        except ImportError:
            # ... otherwise fallback to this code!
            import pymmh3 as mmh3
        # Compute the 128-bit hash from the input string using MurmurHash3
        hash_value = mmh3.hash128(input_string, seed)
        # Extract the low 64 bits
        low_part = hash_value & 0xFFFFFFFFFFFFFFFF
        return low_part

    def set_value(self, input, mask, value):
        shift_count = (mask & -mask).bit_length() - 1
        new_value = (input & ~mask) | ((value << shift_count) & mask)
        return new_value

    def encode_bases_fast(self, alt):
        """Optimized base encoding using dict lookup and direct bit operations"""
        encoded = 0
        for i, c in enumerate(alt):
            value = self.BASE_ENCODING.get(c)
            if value is None:
                raise InvalidPositionException(
                    "Cannot encode alt as it contains strange letters: " + alt
                )
            encoded |= value << (i * 2)
        return encoded

    def hash(self, chromosome, position, del_length, alt):
        change = f"{chromosome}:{position}:{del_length}:{alt}"
        hash_value = self.hash_function(change.encode("utf-8"))
        value = self.set_value(hash_value, self.MASK_HASH, 1)
        return value

    def change_hash(self, del_length, alt):
        change = f"{del_length}:{alt}"
        hash_value = self.hash_function(change.encode("utf-8"))
        value = hash_value & self.MASK_CHANGE_HASH_VALUE
        return value

    def encodeSpdi(self, chromosome: str, position: int, del_length: int, ins: str):
        """
        Encode variant using SPDI format (0-based position, deletion length, insertion string)

        Args:
            chromosome: Chromosome name
            position: 0-based position
            del_length: Length of deletion (refLength)
            ins: Insertion string (normalized alt)
        """
        # Optimized chromosome normalization
        if chromosome.startswith(("chr", "CHR", "Chr")):
            chromosomeNormalized = chromosome[3:].upper()
        else:
            chromosomeNormalized = chromosome.upper()

        positionId = self.positionEncoder.encode_1_based_without_throw(
            chromosomeNormalized, position
        )

        if positionId == -1:
            raise InvalidPositionException(
                "Wrong position: {}:{}.".format(chromosome, position)
            )

        # Fast validation using frozenset
        if not self.VALID_BASES.issuperset(ins):
            raise InvalidPositionException(
                "Alt should contain only ACGNT, position: {}, altNorm: {}.".format(
                    position, ins
                )
            )

        # Check conditions once
        ins_len = len(ins)
        useChangeHash = (
            ins_len > self.MAX_INS_LENGTH or del_length > self.MAX_DEL_LENGTH
        )
        hasN = "N" in ins

        if useChangeHash or hasN:
            # Change hash path
            encoded = 0
            encoded = self.set_value(encoded, self.MASK_C_HASH, 1)
            encoded = self.set_value(encoded, self.MASK_POS, positionId)
            changeHash = self.change_hash(del_length, ins)
            encoded = self.set_value(encoded, self.MASK_CHANGE_HASH_VALUE, changeHash)
            return encoded
        else:
            # Direct encoding path - optimized with direct bit operations
            encodedAlt = self.encode_bases_fast(ins)

            # Combine all values in one operation
            encoded = (
                (positionId << self.SHIFT_POS)
                | (ins_len << self.SHIFT_INS_LEN)
                | (del_length << self.SHIFT_DEL)
                | encodedAlt
            )
            return encoded

    def encode1based(self, chromosome, position, ref, alt):
        # make 0 based
        position = position - 1

        altNorm = alt.upper()

        # Strip leading common characters
        while ref and altNorm and ref[0] == altNorm[0]:
            position += 1
            ref = ref[1:]
            altNorm = altNorm[1:]

        # if ref is string than replace with length
        refLength = len(ref) if isinstance(ref, str) else ref

        # Now call encodeSpdi with the SPDI format parameters
        return self.encodeSpdi(chromosome, position, refLength, altNorm)


# Singleton instance of VariantIdEncoder
_variant_encoder_instance = VariantIdEncoder()


def encode_vcf_variant_gbid(chromosome, position, ref, alt):
    """Optimized function to encode VCF variant"""
    return _variant_encoder_instance.encode1based(chromosome, position, ref, alt)


def encode_spdi_variant_gbid(chromosome, position0, del_length, ins):
    """Optimized function to encode SPDI variant"""
    return _variant_encoder_instance.encodeSpdi(chromosome, position0, del_length, ins)


def encode_positions_vectorized(df_positions):
    """
    Vectorized replacement for:
        df.apply(lambda x: encode_vcf_position_gbid(x["chr"], x["pos"]), axis=1)

    Expects columns:
        chr  – chromosome name (with or without 'chr' / 'CHR' / 'Chr' prefix)
        pos  – 1-based VCF position

    Returns a numpy int64 array of position IDs, one per row.
    Unsupported chromosomes yield PositionEncoder.CHR_NOT_SUPPORTED (-1);
    out-of-range positions yield PositionEncoder.ERROR_WRONG_CHR_POSITION (-2),
    matching the behaviour of the scalar encode_vcf_position_gbid.
    """
    chr_arr = df_positions["chr"].to_numpy()
    pos_arr = np.asarray(df_positions["pos"], dtype=np.int64)
    n = len(df_positions)

    pos_encoder = _position_encoder_instance

    # Build fast lookup arrays (indexed by chromosome index in pos_encoder.chromosomes).
    offsets_arr = np.array(pos_encoder.offsets, dtype=np.int64)
    sizes_arr   = np.array(pos_encoder.sizes,   dtype=np.int64)
    chr_to_idx  = {c: i for i, c in enumerate(pos_encoder.chromosomes)}

    # Strip optional 'chr' prefix and uppercase, matching VariantIdEncoder.encodeSpdi.
    def _norm(c):
        u = str(c).upper()
        return u[3:] if u.startswith("CHR") else u

    normalized = np.array([_norm(c) for c in chr_arr])

    # Default: every row starts as "chromosome not supported".
    result = np.full(n, PositionEncoder.CHR_NOT_SUPPORTED, dtype=np.int64)

    # Process each known chromosome in a single vectorized pass.
    for chr_name, idx in chr_to_idx.items():
        mask = normalized == chr_name
        if not mask.any():
            continue
        p      = pos_arr[mask]
        size   = int(sizes_arr[idx])
        offset = int(offsets_arr[idx])

        # Replicate encode_1_based_without_throw: invalid → sentinel, valid → p-1+offset.
        valid         = (p >= 0) & (p <= size)
        result[mask]  = np.where(
            valid,
            p - 1 + offset,
            PositionEncoder.ERROR_WRONG_CHR_POSITION,
        )

    return result


def encode_spdi_variants_vectorized(df_spdi):
    """
    Vectorized replacement for:
        df.apply(
            lambda x: encode_spdi_variant_gbid(x["chr"], x["pos"], x["del_length"], x["ins"]),
            axis=1,
        )

    Expects columns:
        chr        – chromosome name (with or without 'chr' / 'CHR' / 'Chr' prefix)
        pos        – 0-based SPDI position
        del_length – number of reference bases deleted
        ins        – insertion sequence (may be empty string ""; mixed case is accepted)

    Returns a numpy int64 array of encoded variant IDs, one per row.
    Raises InvalidPositionException for unsupported chromosomes or invalid bases.
    """
    chr_arr     = df_spdi["chr"].to_numpy()
    pos_arr     = np.asarray(df_spdi["pos"], dtype=np.int64)
    del_arr     = np.asarray(df_spdi["del_length"], dtype=np.int64)
    ins_raw     = df_spdi["ins"].to_numpy()

    encoder = _variant_encoder_instance
    n = len(df_spdi)
    result = np.zeros(n, dtype=np.int64)

    # ── 1. Normalize chromosome names ───────────────────────────────────────
    def _norm(c):
        u = str(c).upper()
        return u[3:] if u.startswith("CHR") else u

    chr_norm = np.array([_norm(c) for c in chr_arr])

    # ── 2. Uppercase insertion sequences ────────────────────────────────────
    # encodeSpdi validates against frozenset("ACGNT") which is all uppercase,
    # so normalise here to match scalar behaviour.
    ins_list = [str(a).upper() for a in ins_raw]
    ins_lens = np.array([len(a) for a in ins_list], dtype=np.int64)

    # ── 3. Encode positions in bulk (SPDI 0-based pos passed directly) ──────
    pos_encoder = encoder.positionEncoder
    position_ids = np.array([
        pos_encoder.encode_1_based_without_throw(c, int(p))
        for c, p in zip(chr_norm, pos_arr)
    ])

    # Unsupported chromosome → CHR_NOT_SUPPORTED (-1) → InvalidPositionException.
    invalid_mask = position_ids == PositionEncoder.CHR_NOT_SUPPORTED
    if invalid_mask.any():
        bad = int(np.where(invalid_mask)[0][0])
        raise InvalidPositionException(
            f"Wrong position: {chr_arr[bad]}:{pos_arr[bad]}"
        )

    # ── 4. Classify rows: change-hash path vs direct encoding path ──────────
    has_N    = np.array(["N" in a for a in ins_list])
    use_hash = (ins_lens > encoder.MAX_INS_LENGTH) | (del_arr > encoder.MAX_DEL_LENGTH) | has_N

    # ── 5a. Direct encoding path (fully vectorized bit operations) ───────────
    direct_mask = ~use_hash
    if direct_mask.any():
        d_pos_ids  = position_ids[direct_mask]
        d_ins_lens = ins_lens[direct_mask]
        d_del_lens = del_arr[direct_mask]
        d_ins      = [ins_list[i] for i in np.where(direct_mask)[0]]

        # encode_bases_fast validates bases; will raise on unexpected characters.
        encoded_alts = np.array(
            [encoder.encode_bases_fast(a) for a in d_ins], dtype=np.int64
        )

        result[direct_mask] = (
            (d_pos_ids  << encoder.SHIFT_POS)
            | (d_ins_lens << encoder.SHIFT_INS_LEN)
            | (d_del_lens << encoder.SHIFT_DEL)
            | encoded_alts
        )

    # ── 5b. Change-hash path (unavoidably serial, but typically rare) ────────
    for i in np.where(use_hash)[0]:
        encoded = 0
        encoded = encoder.set_value(encoded, encoder.MASK_C_HASH, 1)
        encoded = encoder.set_value(encoded, encoder.MASK_POS, int(position_ids[i]))
        change_hash = encoder.change_hash(int(del_arr[i]), ins_list[i])
        encoded = encoder.set_value(encoded, encoder.MASK_CHANGE_HASH_VALUE, change_hash)
        result[i] = encoded

    return result


def encode_variants_vectorized(df_variants):
    """
    Vectorized replacement for:
        df_variants["gbid"] = df_variants.apply(
            lambda x: gnb.encode_vcf_variant_gbid(x["chr"], x["pos"], x["ref"], x["alt"]), axis=1
        )
    """
    chr_arr = df_variants["chr"].to_numpy()
    pos_arr = df_variants["pos"].to_numpy()
    ref_arr = df_variants["ref"].to_numpy()
    alt_arr = df_variants["alt"].to_numpy()

    encoder = _variant_encoder_instance

    n = len(df_variants)
    result = np.zeros(n, dtype=np.int64)

    # ── 1. Normalize chromosomes (vectorized string strip) ──────────────────
    def normalize_chr(chrom):
        c = chrom.upper()
        return c[3:] if c.startswith("CHR") else c

    chr_norm = np.array([normalize_chr(c) for c in chr_arr])

    # ── 2. Strip shared leading bases & compute 0-based position ────────────
    pos0 = pos_arr - 1  # convert to 0-based (copy, don't mutate)
    ref_stripped = ref_arr.copy().astype(object)
    alt_stripped = alt_arr.copy().astype(object)

    # Vectorized leading-base strip via numpy where
    # (Loop only over unique (ref,alt) pairs to reduce work)
    for i in range(n):
        r, a = str(ref_stripped[i]).upper(), str(alt_stripped[i]).upper()
        p = pos0[i]
        while r and a and r[0] == a[0]:
            p += 1
            r, a = r[1:], a[1:]
        pos0[i] = p
        ref_stripped[i] = r
        alt_stripped[i] = a

    # ── 3. Encode positions in bulk ─────────────────────────────────────────
    pos_encoder = encoder.positionEncoder
    position_ids = np.array([
        pos_encoder.encode_1_based_without_throw(c, p)
        for c, p in zip(chr_norm, pos0)
    ])

    invalid_mask = position_ids == -1
    if invalid_mask.any():
        bad = np.where(invalid_mask)[0][0]
        raise InvalidPositionException(
            f"Wrong position: {chr_arr[bad]}:{pos0[bad]}"
        )

    # ── 4. Compute del_length and ins arrays ────────────────────────────────
    del_lengths = np.array([len(r) for r in ref_stripped], dtype=np.int64)
    ins_list    = [str(a) for a in alt_stripped]
    ins_lens    = np.array([len(a) for a in ins_list], dtype=np.int64)

    # ── 5. Classify rows: hash path vs direct path ──────────────────────────
    has_N          = np.array(["N" in a for a in ins_list])
    use_hash       = (ins_lens > encoder.MAX_INS_LENGTH) | (del_lengths > encoder.MAX_DEL_LENGTH) | has_N

    # ── 6a. Direct encoding path (vectorized bit ops) ───────────────────────
    direct_mask = ~use_hash
    if direct_mask.any():
        d_pos_ids  = position_ids[direct_mask]
        d_ins_lens = ins_lens[direct_mask]
        d_del_lens = del_lengths[direct_mask]
        d_ins      = [ins_list[i] for i in np.where(direct_mask)[0]]

        encoded_alts = np.array([encoder.encode_bases_fast(a) for a in d_ins], dtype=np.int64)

        result[direct_mask] = (
            (d_pos_ids  << encoder.SHIFT_POS)
            | (d_ins_lens << encoder.SHIFT_INS_LEN)
            | (d_del_lens << encoder.SHIFT_DEL)
            | encoded_alts
        )

    # ── 6b. Hash path (unavoidably serial, but typically rare) ──────────────
    hash_indices = np.where(use_hash)[0]
    for i in hash_indices:
        encoded = 0
        encoded = encoder.set_value(encoded, encoder.MASK_C_HASH, 1)
        encoded = encoder.set_value(encoded, encoder.MASK_POS, int(position_ids[i]))
        change_hash = encoder.change_hash(int(del_lengths[i]), ins_list[i])
        encoded = encoder.set_value(encoded, encoder.MASK_CHANGE_HASH_VALUE, change_hash)
        result[i] = encoded

    return result


# Example usage:
# v = VariantIdEncoder()
# print(v.encode1based("1", 12, "A", "C"))
