#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright 2024 EMBL - European Bioinformatics Institute
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from Bio.Seq import Seq
from pathlib import Path
import numpy as np

from pimento.bin.pimento_utils import (
    get_read_count,
    compute_windowed_base_conservation,
    build_list_of_base_counts,
    fetch_read_substrings,
)
from pimento.bin.thresholds import CONSENSUS_BASE_THRESHOLD, MAX_READ_COUNT


def choose_cutoff_for_single_strand(
    input_fastq: Path, cutoff_list: list[int], rev: bool = False
) -> None:
    """
    Assess inflection point list, selecting one for automatic primer trimming.

    Takes as input a fastq file and a list of inflection points generated by "find_cutoffs.py".
    Computes the average conservation of bases before inflection point and after.
    Gets the difference in avg. conservation between the pre- and post- points.
    Selects the inflection point with the maximum difference as the cutoff.
    If an inf point has a similar difference and is earlier than the max, we make a 'conservative' choice and
    replace it with the earlier cutoff.

    Returns the cutoff point and the consensus sequence 'forming' the automatically predicted primer
    """

    start_confs = []  # pre-inf point conservations
    end_confs = []  # post-inf point conservations
    start_cons_lens = []  # list for storing lengths of pre-inflection point sequences
    cons_seq_list = []  # list for storing consensus sequences pre-inflection points

    do_not_include_list = [
        i + 5 for i in cutoff_list
    ]  # ignore conservation of inflection point in calculation

    read_count = get_read_count(input_fastq)  # get readcount from fastq

    max_line_count = 0
    if read_count > MAX_READ_COUNT:
        max_line_count = MAX_READ_COUNT

    for start in cutoff_list:  # Looping through the pre-inflection point substrings
        substring_len = (
            start + 4
        )  # length of pre-inf substrings is inflection point + 4

        read_substring_count_dict = fetch_read_substrings(
            input_fastq, substring_len, rev=rev, max_line_count=max_line_count
        )  # get base count dict
        base_counts = build_list_of_base_counts(
            read_substring_count_dict, substring_len
        )  # list of base conservation dicts for substrings
        base_conservation, cons_seq = compute_windowed_base_conservation(
            base_counts,
            read_count,
            CONSENSUS_BASE_THRESHOLD,
            do_not_include_list,
            max_line_count=max_line_count,
        )  # get list of max base conservations for each index
        # also get consensus sequence
        cons_seq_list.append(cons_seq)
        start_confs.append(np.mean(base_conservation))
        start_cons_lens.append(len(cons_seq))

    for i, end in enumerate(
        cutoff_list
    ):  # Looping through the post-inflection point substrings
        substring_len = end + 5  # length of pre-inf substrings is inflection point + 5
        subs_len = start_cons_lens[i]  # length of respective pre-inf point sequence

        read_substring_count_dict = fetch_read_substrings(
            input_fastq, subs_len, rev, substring_len, max_line_count=max_line_count
        )
        base_counts = build_list_of_base_counts(read_substring_count_dict, subs_len)
        base_conservation, cons_seq = compute_windowed_base_conservation(
            base_counts,
            read_count,
            CONSENSUS_BASE_THRESHOLD,
            do_not_include_list,
            subs_len,
            max_line_count=max_line_count,
        )

        end_confs.append(np.mean(base_conservation))

    diff_res = [
        start_confs[i] - end_confs[i] for i in range(len(start_confs))
    ]  # get differences between pre- and -post avg conservation values
    diff_res_sorted = sorted(
        diff_res, reverse=True
    )  # sort differences from highest to lowest

    ini_max_res = diff_res_sorted[0]  # maximum differences
    curr_max_index = diff_res.index(ini_max_res)  # index of maximum differences

    for res in diff_res_sorted[1:]:  # Loop through the rest of the differences
        curr_res_index = np.where(diff_res == res)[0][0]

        index_diff = cutoff_list[curr_max_index] - cutoff_list[curr_res_index]

        # if difference between the max and the current is negligible and the index of the current is earlier then..
        if ini_max_res - res < 0.05 and (index_diff <= 3 and index_diff > 0):
            curr_max_index = (
                curr_res_index  # replace the selected index with the current one
            )

    cutoff = cutoff_list[curr_max_index] + 5  # cutoff is the inflection point index + 5
    primer = cons_seq_list[
        curr_max_index
    ]  # grab the correct consensus sequence as primer

    # if the requested strand is reverse..
    if rev:
        primer = str(
            Seq(primer).complement()
        )  # ..get the complement of consensus sequence

    return cutoff, primer
