#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import matplotlib.patches as mp
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import pandas as pd
import pysam
import seaborn as sns
from typing import Optional, Sequence, Tuple, Mapping
import pyBigWig

######
# plot gene model
######

def plot_gene_model(
    ax,
    gene_models: pd.DataFrame,
    fig_start: int,
    fig_end: int,
    gene_color: str = "k",
    y_space: int = 1,
    plot_gene_id: bool = True
):
    """plot gene model in the axis

    Args:
    -----
        ax (matplotlib.axes): An axis object to plot.

        gene_models (pd.DataFrame): A bed12 like DataFrame:

            chrom     start       end      gene_id score strand thickStart  thickEnd
                2  17922018  17924542  AT2G43110.1     0      +   17922066  17924291

            rgb blockCount              blockSizes                blockStarts y_pos
              0          6  361,196,73,50,114,372,  0,527,823,1802,1952,2152,     0

        fig_start (int): figure xlim start.

        fig_end (int): figure xlim end.

        gene_color (str, optional): gene color. Defaults to 'k'.

        y_space (int, optional): the spaces between gene in y direction. Defaults to 1.
    """
    ylim = 0  # ax ylim的下限
    height = 3  # gene model 高度
    y_space = y_space + height * 2
    for gene_model in gene_models.values:
        (
            chrom,
            chromStart,
            chromEnd,
            gene_id,
            _,
            strand,
            thickStart,
            thickEnd,
            _,
            blockCount,
            blockSizes,
            blockStarts,
            y_pos,
        ) = gene_model
        y_pos = -y_space * y_pos
        ylim = min(y_pos, ylim)

        # 数据类型转化
        chromStart = int(chromStart)
        chromEnd = int(chromEnd)
        thickStart = int(thickStart)
        thickEnd = int(thickEnd)
        blockSizes = np.fromstring(blockSizes, sep=",", dtype="int")
        blockStarts = np.fromstring(blockStarts, sep=",", dtype="int") + chromStart

        # 画转录起始位点及方向箭头
        small_relative = 0.06 * (fig_end - fig_start)  # 箭头突出部分相对长度
        arrowprops = dict(arrowstyle="-|>", connectionstyle="angle", color=gene_color)
        if strand == "+":
            ax.annotate(
                "",
                xy=(chromStart + small_relative, height * 2 + y_pos),
                xytext=(chromStart, y_pos),
                arrowprops=arrowprops,
            )
        else:
            ax.annotate(
                "",
                xy=(chromEnd - small_relative, height * 2 + y_pos),
                xytext=(chromEnd, y_pos),
                arrowprops=arrowprops,
            )

        line = mp.Rectangle(
            (chromStart, y_pos - height / 8),
            chromEnd - chromStart,
            height / 4,
            color=gene_color,
            linewidth=0,
        )  # 基因有多长这条线就有多长
        ax.add_patch(line)

        # draw exons
        cds_split = [thickStart]
        for exonstart, size in zip(blockStarts, blockSizes):
            exon = mp.Rectangle(
                    (exonstart, y_pos - height / 2),
                    size,
                    height,
                    color=gene_color,
                    linewidth=0,
                )
            ax.add_patch(exon)
            for exon_edge in [exonstart, exonstart + size]:
                if thickStart < exon_edge and exon_edge < thickEnd:
                    cds_split.append(exon_edge)
        cds_split.append(thickEnd)

        thick_list = [(cds_split[i], cds_split[i+1]) for i in range(0, len(cds_split), 2)]

        for thick_pair in thick_list:
            thick = mp.Rectangle((thick_pair[0], y_pos-height), thick_pair[1]-thick_pair[0], height*2, color=gene_color, linewidth=0)
            ax.add_patch(thick)

        # for exonstart, size in zip(blockStarts, blockSizes):
        #     if exonstart == chromStart and exonstart + size == chromEnd:
        #         utr_size = thickStart - chromStart
        #         utr = mp.Rectangle(
        #             (exonstart, y_pos - height / 2),
        #             utr_size,
        #             height,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(utr)
        #         utr_size = chromEnd - thickEnd
        #         utr = mp.Rectangle(
        #             (thickEnd, y_pos - height / 2),
        #             utr_size,
        #             height,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(utr)
        #         exon = mp.Rectangle(
        #             (thickStart, y_pos - height),
        #             thickEnd - thickStart,
        #             height * 2,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(exon)

        #     elif exonstart + size <= thickStart:
        #         # 只有5'/ 3'UTR
        #         utr = mp.Rectangle(
        #             (exonstart, y_pos - height / 2),
        #             size,
        #             height,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(utr)

        #     elif exonstart < thickStart and exonstart + size > thickEnd:
        #         # cds surrounded by utrs
        #         utr = mp.Rectangle((exonstart, y_pos-height/2), size, height, color=gene_color, linewidth=0)
        #         exon = mp.Rectangle((thickStart, y_pos-height), thickEnd-thickStart, height*2, color=gene_color, linewidth=0)
        #         ax.add_patch(utr)
        #         ax.add_patch(exon)

        #     elif exonstart < thickStart and exonstart + size > thickStart:
        #         # 带有5' / 3' UTR的exon
        #         utr_size = thickStart - exonstart
        #         utr = mp.Rectangle(
        #             (exonstart, y_pos - height / 2),
        #             utr_size,
        #             height,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         exon = mp.Rectangle(
        #             (exonstart + utr_size, y_pos - height),
        #             size - utr_size,
        #             height * 2,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(utr)
        #         ax.add_patch(exon)

        #     elif exonstart >= thickStart and exonstart + size <= thickEnd:
        #         # 普通exon
        #         exon = mp.Rectangle(
        #             (exonstart, y_pos - height),
        #             size,
        #             height * 2,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(exon)

        #     elif exonstart < thickEnd and exonstart + size > thickEnd:
        #         # 带有3' / 5' UTR的exon
        #         utr_size = exonstart + size - thickEnd
        #         utr = mp.Rectangle(
        #             (thickEnd, y_pos - height / 2),
        #             utr_size,
        #             height,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         exon = mp.Rectangle(
        #             (exonstart, y_pos - height),
        #             size - utr_size,
        #             height * 2,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(utr)
        #         ax.add_patch(exon)

        #     elif exonstart >= thickEnd:
        #         # 只有3'/ 5'UTR
        #         utr = mp.Rectangle(
        #             (exonstart, y_pos - height / 2),
        #             size,
        #             height,
        #             color=gene_color,
        #             linewidth=0,
        #         )
        #         ax.add_patch(utr)

        if plot_gene_id:
            ax.annotate(
                gene_id, xy=((chromStart + chromEnd) / 2, y_pos + height * 1.5), ha="center"
            )

    # set ax
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.spines["bottom"].set_visible(False)
    ax.yaxis.set_major_locator(ticker.NullLocator())
    ax.xaxis.set_major_locator(ticker.NullLocator())
    ax.xaxis.set_ticks_position("none")

    ax.set_ylim(ylim - height, height + y_space)
    ax.set_xlim(fig_start, fig_end)

def get_gene_model(chrom: str, start: int, end: int, bed_path: str, lambda_for_parse_geneId:Optional[str] = None) -> pd.DataFrame:
    """Get gene model information from the bed file.
        The bed file should be indexed by tabix.
        For details, see http://www.htslib.org/doc/tabix.html

    Args:
        chrom (str): chromosome id.

        start (int): the start of the region.

        end (int): the end of the region.

        bed_path (str): the PATH of the bed file.

        lambda_for_parse_geneId(str) : lambda expression for parse geneId. Defalut is x: x

    Returns:
        pd.DataFrame: A bed12 dataframe.
            Examples
            --------
            chrom     start       end      gene_id score strand thickStart  thickEnd
                2  17922018  17924542  AT2G43110.1     0      +   17922066  17924291

            rgb blockCount              blockSizes                blockStarts
              0          6  361,196,73,50,114,372,  0,527,823,1802,1952,2152,
    """
    tbx = pysam.TabixFile(bed_path)
    gene_models = pd.DataFrame(
        [gene_model.split("\t") for gene_model in tbx.fetch(chrom, start, end)],
        columns=[
            "chrom",
            "start",
            "end",
            "gene_id",
            "score",
            "strand",
            "thickStart",
            "thickEnd",
            "rgb",
            "blockCount",
            "blockSizes",
            "blockStarts",
        ],
    )
    gene_models.sort_values(["start", "end"], inplace=True)
    if lambda_for_parse_geneId:
        lambda_for_parse_geneId_fc = eval(f"lambda {lambda_for_parse_geneId}")
        gene_models = gene_models.assign(gene_id = lambda df:df['gene_id'].map(lambda_for_parse_geneId_fc))
    return gene_models

######
# other function
######

def is_overlap(
    gene_a: Tuple[str, int, int], gene_b: Tuple[str, int, int], threshold: int = 0
) -> bool:
    """To judge whether two region is overlap.

    Args:
        gene_a (tuple): (chrom, start, end)
        gene_b (tuple): (chrom, start, end)
        threshold (int, optional): the minimum space between two region. Defaults to 0.

    Returns:
        bool: if overlap True, else False
    """
    minn = max(int(gene_a[1]), int(gene_b[1]))
    maxn = min(int(gene_a[2]), int(gene_b[2]))

    if maxn - minn >= -threshold:
        return True
    else:
        return False

def get_y_pos_discontinous(df: pd.DataFrame, gene_list: Optional[Sequence[str]] = None):
    """Get the y position of each region. Save the results in the df['y_pos'].

    Args:
        df (pd.DataFrame): a DataFrame must include first four columns:
            chrom, start, end, gene_id.

            Examples
            --------
            chrom    start      end    gene_id strand
                4  9105672  9106504  AT4G16100      +

        gene_list (set, optional): a set contain which gene to plot.
            When gene_list is None, plot all item in the df. Defaults to None.

    Returns:
        int: item counts in the y position.
    """

    # only plot reads/gene_model in the gene_list
    if gene_list is not None:
        df.query("gene_id in @gene_list", inplace=True)

    df.reset_index(drop=True, inplace=True)
    df["y_pos"] = df.index

    return len(df)

def filter_bam(
    df,
    strand=None,
    start_before=None,
    start_after=None,
    end_before=None,
    end_after=None,
):
    if strand is not None:
        df.query("strand == @strand", inplace=True)
    if start_before is not None:
        df.query("start <= @start_before", inplace=True)
    if start_after is not None:
        df.query("start >= @start_after", inplace=True)
    if end_before is not None:
        df.query("end <= @end_before", inplace=True)
    if end_after is not None:
        df.query("end >= @end_after", inplace=True)

def get_y_pos_continuous(df, gene_list=None, threshold=0):
    """Get the y position of each region. Save the results in the df['y_pos'].

    Args:
        df (pd.DataFrame): a DataFrame must include first four columns:
            chrom, start, end, gene_id.

            Examples
            --------
            chrom    start      end    gene_id strand
                4  9105672  9106504  AT4G16100      +

        gene_list (set, optional): a set contain which gene to plot.
            When gene_list is None, plot all item in the df. Defaults to None.

        threshold (int, optional): the minimum space between two region. Defaults to 0.

    Returns:
        int: item counts in the y position.
    """
    # initialization of y_pos columns
    df["y_pos"] = None

    read_list = []
    for index_id, item in enumerate(df.values):
        chrom, start, end, gene_id, *_ = item

        # only plot reads/gene_model in the gene_list
        if gene_list is not None and gene_id not in gene_list:
            df.drop(index_id, inplace=True)
            continue

        current_read = (chrom, start, end)
        is_add_read = False
        y_pos = -1
        for y_pos in range(len(read_list)):
            y_pos_read = read_list[y_pos]
            if not is_overlap(current_read, y_pos_read, threshold=threshold):
                read_list[y_pos] = current_read
                df.at[index_id, "y_pos"] = y_pos
                is_add_read = True
                break
        if not is_add_read:
            read_list.append(current_read)
            y_pos += 1
            df.at[index_id, "y_pos"] = y_pos

    return len(read_list)

def get_splice_sites(x):
    splice_sites = []
    for i in x:
        splice_sites.append(i[0])
        splice_sites.append(i[0]+i[1])

    return tuple(splice_sites[1:])

######
# plot bam
######

def find_exon(read):
    """Find exon position in the pysam.read.

    Args:
        read (pysam.read): <pysam.libcalignedsegment.AlignedSegment>

    Returns:
        zip(exon_start, exon_size): the exon position in the read.
            exon_start: exon start coordinate, 0-based
            exon_size: exon length

        Examples
        --------
        [(9105672, 457), (9106400, 142)]
    """
    BAM_CREF_SKIP = 3  # BAM_CREF_SKIP
    blockStart = []
    blockSize = []
    match_or_deletion = {
        0,
        2,
        7,
        8,
    }  # only M/=/X (0/7/8) and D (2) are related to genome position
    exon_start = read.reference_start
    length = 0
    for op, nt in read.cigartuples:
        if op in match_or_deletion:
            length += nt
        elif op == BAM_CREF_SKIP:
            blockStart.append(exon_start)
            blockSize.append(length)
            exon_start += length + nt
            length = 0
    blockStart.append(exon_start)
    blockSize.append(length)
    return zip(blockStart, blockSize)

def convert_bam(
    chrom,
    start,
    end,
    strand,
    infile,
    subsample=None,
    gene_list=None,
    method="continuous",
    filter_strand=None,
    start_before=None,
    start_after=None,
    end_before=None,
    end_after=None,
    force_tag_check: bool = True,
    tag_required_dict:Mapping[str, str]=dict(gene_id='gi', polya_len='pa', span_intron_count='sn', unsplice_count='rn', unsplice_intron='ri'),
    chrom_prefix='',
    color='#5D93C4',
):
    """Conver pysam.reads into a DataFrame.

    Args:
    -----

        chrom (str): chromosome id

        start (int): start position of the region

        end (int): end position of the region

        strand (str): '+' or '-'

        infile (str): the PATH of the bam file

        subsample (int or float, optional): <number>|<frac> of items from axis to return.
            Defaults to None.

        gene_list (set, optional): a set contain which gene to plot.
            When gene_list is None, plot all item in the df. Defaults to None.

        method ('continuous' | '3_end' | '5_end')

    Returns:
        DataFrame: A dataframe contain bam reads information
            Examples
            --------
            Return: pd.DataFrame
            chrom    start      end    gene_id strand
                4  9105672  9106504  AT4G16100      +

                                         read_id  gap  polya_len
            37c91155-0ef5-47be-9e1b-d3a6753982f4  0.0          0

                        exon     y_pos
            [(9105672, 832)]      None
    """

    bam_data = []
    chrom = chrom_prefix+chrom
    with pysam.AlignmentFile(infile, "rb") as inbam:
        for read in inbam.fetch(chrom, start, end):
            gap = 0  # useless value

            try:
                polya_len = read.get_tag('pa')
            except KeyError:
                polya_len = 0

            try:
                gene_id = read.get_tag('gi')
            except KeyError:
                gene_id = 0

            try:
                span_intron_count = read.get_tag('sn')  # span_intron_num
                if type(span_intron_count) == str:
                    span_intron_count = len(span_intron_count.split(':'))
                unsplice_count = read.get_tag('rn')  # retention_intron_num
                unsplice_intron = read.get_tag('ri')  # retention_introns
                # only plot reads in gene_list
            except KeyError:
                span_intron_count, unsplice_count, unsplice_intron = None, None, None

            '''新增标记是否为novel isoform的tag，0为False，1为True'''
            try:
                novel_isoform = read.get_tag('ni')
            except KeyError:
                novel_isoform = 0

            '''新增标记是否chimeric read的tag，0为False，1为True'''
            try:
                chimeric_read = read.get_tag('cr')
            except KeyError:
                chimeric_read = 0

            if read.is_supplementary or read.is_unmapped:
                continue

            read_strand = "-" if read.is_reverse else "+"
            exon = find_exon(read)

            bam_data.append(
                (
                    read.reference_name,
                    read.reference_start,
                    read.reference_end,
                    gene_id,
                    read_strand,
                    read.query_name,
                    gap,
                    polya_len,
                    list(exon),
                    span_intron_count,
                    unsplice_count,
                    unsplice_intron,
                    novel_isoform,
                    chimeric_read,
                )
            )
    bam_data = pd.DataFrame(
        bam_data,
        columns=[
            "chrom",
            "start",
            "end",
            "gene_id",
            "strand",
            "read_id",
            "gap",
            "polya_len",
            "exon",
            "span_intron_count",
            "unsplice_count",
            "unsplice_intron",
            "novel_isoform",
            "chimeric_read",
        ],
    )

    bam_data['color'] = color  # read color

    # 进行subsample
    if subsample is not None:
        if type(subsample) is int:
            bam_data = bam_data.sample(n=subsample, random_state=42)
        elif type(subsample) is float and subsample < 1:
            bam_data = bam_data.sample(frac=subsample, random_state=42)

    if method == "continuous":
        # 类似于igv那样连续的排
        # add polya
        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )

        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )
        if strand == "+":
            bam_data.sort_values(["start", "end"], inplace=True)
        else:
            bam_data.sort_values(["end", "start"], ascending=False, inplace=True)
        bam_data.reset_index(drop=True, inplace=True)
        get_y_pos_continuous(bam_data, gene_list=gene_list)

    elif method == "5_end":
        if strand == "+":
            bam_data.sort_values(["start", "end"], inplace=True)
        else:
            bam_data.sort_values(["end", "start"], ascending=False, inplace=True)

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )

        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    elif method == "3_end":
        if strand == "+":
            bam_data.sort_values(["end", "start"], inplace=True)
        else:
            bam_data.sort_values(["start", "end"], ascending=False, inplace=True)

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )

        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    elif method == "gene_id":  # 按照gene_id 3' end排序
        if strand == "+":
            bam_data.sort_values(["gene_id", "end", "start"], inplace=True)
        else:
            bam_data.sort_values(
                ["gene_id", "start", "end"], ascending=False, inplace=True
            )

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    elif method == "spliced":  # 只画完全剪切的reads
        bam_data.query("span_intron_count > 0 and unsplice_count == 0", inplace=True)
        if strand == "+":
            bam_data.sort_values(["end", "start"], inplace=True)
        else:
            bam_data.sort_values(["start", "end"], ascending=False, inplace=True)

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )
        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    elif method == "partially_spliced":  # 只画不完全剪切的reads
        bam_data.query(
            "span_intron_count > 0 and unsplice_count > 0 and unsplice_count != span_intron_count ",
            inplace=True,
        )
        if strand == "+":
            bam_data.sort_values(["end", "start"], inplace=True)
        else:
            bam_data.sort_values(["start", "end"], ascending=False, inplace=True)

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )
        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    elif method == "unspliced":  # 只画不完全剪切的reads
        bam_data.query(
            "span_intron_count > 0 and unsplice_count == span_intron_count",
            inplace=True,
        )
        if strand == "+":
            bam_data.sort_values(["end", "start"], inplace=True)
        else:
            bam_data.sort_values(["start", "end"], ascending=False, inplace=True)

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )
        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    elif method == "ir":
        bam_data['exon1'] = bam_data['exon'].map(get_splice_sites)
        if strand == "+":
            bam_data.sort_values(['novel_isoform', 'exon1', "end", "start"], inplace=True)
        else:
            bam_data.sort_values(['exon1', "start", "end"], ascending=False, inplace=True)
            bam_data.sort_values(['novel_isoform'], inplace=True)

        bam_data['color'] = bam_data.apply(lambda x: '#A6CEE3' if x.novel_isoform == 0 else x.color, axis=1)

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )

        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    elif method == "chimeric":
        if strand == "+":
            bam_data.sort_values(["chimeric_read", "end", "start"], inplace=True)
        else:
            bam_data['chimeric_read'] = -1 * bam_data['chimeric_read']
            bam_data.sort_values(["chimeric_read", 'start', 'end'], ascending=False, inplace=True)

        bam_data['color'] = bam_data.apply(lambda x: '#A6CEE3' if x.chimeric_read == 0 else x.color, axis=1)

        # add polya
        bam_data["polya_len"] = bam_data["polya_len"].map(
            lambda x: int(x) if x >= 15 else 0
        )
        bam_data = bam_data.assign(
            start=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["start"], bam_data["start"] - bam_data["polya_len"]],
                -1,
            ),
            end=np.select(
                [bam_data["strand"] == "+", bam_data["strand"] == "-"],
                [bam_data["end"] + bam_data["polya_len"], bam_data["end"]],
                -1,
            ),
        )

        filter_bam(
            bam_data,
            strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
        )
        get_y_pos_discontinous(bam_data, gene_list=gene_list)

    return bam_data

def plot_bam(
    ax,
    bam_data,
    fig_start,
    fig_end,
    read_color="#5D93C4",  # 已经无用
    polya_color="lightcoral",
    y_space=1.5,  # the space between reads in yaxis
    gene_list=None,
    pal=None,
    height: int = 0.5, # reads的高度
):
    """plot bam information to the ax

    Args:
    -----
        ax (matplotlib.axes): An axis object to plot

        bam_data (pd.DataFrame): A dataframe contain bam reads information
            Examples
            --------
            Return: pd.DataFrame
            chrom    start      end    gene_id strand
                4  9105672  9106504  AT4G16100      +

                                         read_id  gap  polya_len
            37c91155-0ef5-47be-9e1b-d3a6753982f4  0.0          0

                        exon     y_pos
            [(9105672, 832)]         0

        fig_start (int): figure xlim start.

        fig_end (int): figure xlim end.

        read_color (str, optional): read color. Defaults to '#5D93C4'.

        polya_color (str, optional): polya color. Defaults to 'lightcoral'.

        y_space (int, optional): the spaces between reads in y direction. Defaults to 1.5.

        pal (seaborn.palettes._ColorPalette, optional)
    """
    if pal is None:
        pal = sns.color_palette("Paired")

    gene_color = {}
    gene_color_index = 0

    ylim = 0  # the start position of yaxis
    for item in bam_data.itertuples():
        (
            chrom,
            start,
            end,
            gene_id,
            strand,
            read_id,
            gap,
            polya_len,
            exon,
            *_,
            ypos,
        ) = item[1:]
        ypos = -y_space * ypos
        ylim = min(ypos, ylim)
        line = mp.Rectangle(
            (start, ypos - height / 4),
            end - start,
            height / 2,
            color="#A6A6A6",
            linewidth=0,
        )
        ax.add_patch(line)
        for block_start, block_size in exon:

            # set gene color
            # index为色板中的序列编号
            # 如果gene_list存在 则给gene_list里面都基因上色
            if gene_list is not None:
                if gene_id in gene_list:
                    if gene_id not in gene_color:
                        gene_color[gene_id] = gene_color_index
                        gene_color_index += 1
                    read_color_ = pal[gene_color[gene_id]]
                else:
                    # 不在gene_list里面的reads都设置成灰色
                    read_color_ = "#5D5D5D"
            # 如果没有则统一颜色
            else:
                read_color_ = item.color

                # TODO: 不同链的reads不同颜色
            exon = mp.Rectangle(
                (block_start, ypos - height),
                block_size,
                height * 2,
                color=read_color_,
                linewidth=0,
            )
            ax.add_patch(exon)

        # plot polya
        if polya_len > 15:
            if strand == "+":
                polya_tail = mp.Rectangle(
                    (block_start + block_size, ypos - height),
                    polya_len,
                    height * 2,
                    color=polya_color,
                    linewidth=0,
                )
            else:
                polya_tail = mp.Rectangle(
                    (start, ypos - height),
                    polya_len,
                    height * 2,
                    color=polya_color,
                    linewidth=0,
                )
            ax.add_patch(polya_tail)

    ax.set_ylim(ylim * 1.1, height + y_space)
    # ax.set_xlim(fig_start, fig_end)

def set_ax(ax, plot_xaxis=False):
    """Set axes

    Args:
        ax (matplotlib.axes): An axis object to plot

        plot_xaxis (bool, optional): if True plot x_ticks and x_locator.
            Defaults to False.
    """
    ax.spines["right"].set_visible(False)  # 去线
    ax.spines["left"].set_visible(False)  # 去线
    ax.spines["top"].set_visible(False)  # 去线
    ax.yaxis.set_major_locator(ticker.NullLocator())  # 去y数字

    if not plot_xaxis:
        ax.xaxis.set_major_locator(ticker.NullLocator())  # 去x数字
        ax.xaxis.set_ticks_position("none")  # 去x刻度

def plot_bw_track(ax, track_data, start, end, track_color: str = None, data_range: tuple = (None, None)):
    '''This function takes in a matplotlib axis, a dataframe of track data, a start and end time, and a
    color, and plots the track on the axis

    Parameters
    ----------
    ax
        the axis to plot the track on
    track_data
        a list of tuples, each tuple is a point in the track
    start
        the start time of the track
    end
        the end of the track
    track_color
        the color of the track

    '''
    ax.fill_between(np.linspace(start, end, end-start), track_data, color=track_color)
    # set ax
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.spines["bottom"].set_visible(False)
    ax.yaxis.set_major_locator(ticker.NullLocator())
    ax.xaxis.set_major_locator(ticker.NullLocator())
    ax.xaxis.set_ticks_position("none")
    # ax.set_xlim(start, end)
    if not(None in data_range):
        ax.set_ylim(data_range)

######
# main Class
######

class IGV:
    """A IGV class

    Usage:
    ------
        chrom, start, end, strand, force_tag_check = '5', 5371627, 5375616, '+', True

        igv_plot = igv.IGV(chrom, start, end, strand=strand)

        araport11_isoform_path = '/public/home/mowp/db/Arabidopsis_thaliana/representative_gene_model/araport11.representative.gene_model.bed.gz'
        igv_plot.add_gene_model(
            araport11_isoform_path,
            gene_list = {'AT5G16440.1', 'AT5G16450.1'},
        )

        infile = '/public/home/mowp/workspace/termination/cbRNA_pool/elongating_data/elongating_data.bam'
        igv_plot.add_bam(
            infile,
            gene_list = {'AT5G16440', 'AT5G16450'},
            read_color='r',
            bam_title='one bam'
        )

        igv_plot.plot(height=4, width=8)
    ------
    tag_required_dict:Mapping[str, str]:
        tag required. Key is the variable name, and value is corresponding tag.
        if force_tag_check is True, all tag should be stored in bam; if not, will assign values to variables in order, and when error is triggered, left variables will be ignored.
    """

    def __init__(
        self,
        chrom,
        start,
        end,
        strand="+",
        tag_required_dict:Mapping[str, str]=dict(gene_id='gi', polya_len='pa', span_intron_count='sn', unsplice_count='rn', unsplice_intron='ri'),
    ):
        self.chrom = chrom
        self.start = start
        self.end = end
        self.strand = strand
        self.bam_list = []
        self.tagRequiered = tag_required_dict
        self.colorList = []
        self.bamTitleList = []
        self.gene_models = None
        self.gene_model_ylim = 0
        self.bw_track_list = []

        self.gene_model_n = 0
        self.bw_tracks_n = 0
        self.bam_list_n = 0

    def add_bam(
        self,
        bam_path,
        gene_list=None,
        subsample=None,
        method="continuous",
        filter_strand=None,
        start_before=None,
        start_after=None,
        end_before=None,
        end_after=None,
        read_color='#5D93C4',
        bam_title='',
        chrom_prefix='',
        force_tag_check: bool = True,
    ):

        '''
        chrom_prefix:
            the chromosome name, such as "chr", "Chr", default is ""
        force_tag_check:
            if True, reads which don't have required tags will be ignored
        '''

        bam_data = convert_bam(
            self.chrom,
            self.start,
            self.end,
            self.strand,
            bam_path,
            subsample=subsample,
            gene_list=gene_list,
            method=method,
            filter_strand=filter_strand,
            start_before=start_before,
            start_after=start_after,
            end_before=end_before,
            end_after=end_after,
            force_tag_check=force_tag_check,
            tag_required_dict=self.tagRequiered,
            chrom_prefix=chrom_prefix,
        )
        self.bam_list.append(bam_data)
        self.gene_list = gene_list
        self.colorList.append(read_color)
        self.bamTitleList.append(bam_title)
        self.bam_list_n += 1

    def add_bw(
        self,
        infile: str,
        color='#5D93C4',
        chrom_prefix='',
        data_range: tuple = (None, None),
    ):
        bw = pyBigWig.open(infile)
        chrom = chrom_prefix + self.chrom
        bw_data = bw.values(chrom, self.start, self.end)
        bw_data = np.nan_to_num(bw_data)

        self.bw_track_list.append((bw_data, color, data_range))
        self.bw_tracks_n += 1

    def add_gene_model(self, anno:str, gene_list=None, auto_create_gz=False, lambda_for_parse_geneId:Optional[str]=None, plot_gene_id: bool = True):
        """
        auto_create_gz: if provided bed is not compressed, the bgzip will be used to created the gz format file, corresponding index will also be built.
        lambda_for_parse_geneId: If not None, it will be used to parse the `gene_id` column in bed. None is equal to `x:x`
        """
        if anno.endswith('.gz'):
            self.anno = anno
        else:
            self.anno = anno + '.gz'
            if not os.path.exists(self.anno):
                assert auto_create_gz, "WRONG suffix, and <auto_create_gz> is False"
                print(f"WARNING: Try to create file {self.anno} and build index")
                os.system(f'sort -k1,1 -k2,2n {anno} > temp____ && bgzip -c  temp____ > {self.anno} && tabix -p bed {self.anno}')
        self.gene_models = get_gene_model(self.chrom, self.start, self.end, self.anno, lambda_for_parse_geneId=lambda_for_parse_geneId)
        self.gene_model_ylim = get_y_pos_continuous(
            self.gene_models, gene_list=gene_list, threshold=8
        )  # 不重叠y轴的基因数目
        self.plot_gene_id = plot_gene_id
        self.gene_model_n = 1

    def plot(
        self,
        height=5,
        width=10,
        gene_track_height=0.5,
        bw_track_height=1,
        bam_track_height=4,
        gene_color="k",
        extend_xlim_start=False,
        extend_xlim_end=False,
        polya_site=None,
        read_height: int = 0.5,
        relative_pos: bool = True,
    ):
        '''It plots the data.

        Parameters
        ----------
        height, optional
            the height of the plot
        width, optional
            the width of the plot
        gene_track_height
            height of the gene track
        bw_track_height, optional
            height of the bigwig track
        bam_track_height, optional
            the height of the bam track
        gene_color, optional
            color of the gene track
        extend_xlim_start, optional
            If True, the x-axis will be extended to the left by the length of the left-most read.
        extend_xlim_end, optional
            If True, the x-axis will be extended to the right by the length of the right-most read.
        polya_site
            the position of the polyA site. If not provided, it will be plot the polyA site.
        read_height : int
            the height of the reads in the bam track
        relative_pos : bool, optional
            bool = True

        '''

        gridspec_kw = [gene_track_height]*self.gene_model_n + [bw_track_height]*self.bw_tracks_n + [bam_track_height]*self.bam_list_n
        nrows = self.gene_model_n + self.bw_tracks_n + self.bam_list_n

        fig, axes = plt.subplots(
            figsize=(width, height),
            nrows=nrows,
            gridspec_kw={"height_ratios": gridspec_kw},
            sharex=True,
        )

        if nrows == 1:
            axes = [axes]

        axes_index = 0
        # plot gene_model
        if self.gene_models is not None:
            ax = axes[0]
            plot_gene_model(
                ax,
                self.gene_models,
                self.start,
                self.end,
                gene_color=gene_color,
                y_space=6,
                plot_gene_id=self.plot_gene_id,
            )
            ax.patch.set_alpha(0)
            axes_index += 1

        # plot bw_tracks
        for bw_data, color, data_range in self.bw_track_list:
            plot_bw_track(axes[axes_index], bw_data, self.start, self.end, color, data_range)
            axes[axes_index].patch.set_alpha(0)
            axes_index += 1

        # plot bam files
        for bam_data, read_color, bamTitle in zip(self.bam_list, self.colorList, self.bamTitleList):
            ax = axes[axes_index]

            plot_bam(ax, bam_data, self.start, self.end, gene_list=self.gene_list, read_color=read_color, height=read_height)
            plot_xaxis = axes_index == nrows - 1
            set_ax(ax, plot_xaxis=plot_xaxis)

            # add polya site
            if polya_site is not None:
                tbx = pysam.TabixFile(polya_site)
                for item in tbx.fetch(self.chrom, self.start, self.end):
                    chrom_, start_, end_, gene_id_, _, strand_, _ = item.split("\t")
                    gene_id_ = gene_id_.split("_")[0]
                    if self.gene_list is not None:
                        if gene_id_ in self.gene_list:
                            ax.axvline(int(end_), ls="--", color="#555555")
                    else:
                        ax.axvline(int(end_), ls="--", color="#555555")
            plt.sca(ax)
            plt.text(1.0, 0.5, bamTitle, transform=ax.transAxes,
                fontsize=12, va='center', ha='left')
            ax.patch.set_alpha(0)
            axes_index += 1

        if relative_pos:
            # set last xaxis
            maxn = self.end
            minn = self.start
            for bam_data in self.bam_list:
                if len(bam_data) == 0:
                    continue

                minn_ = min(bam_data["start"])
                maxn_ = max(bam_data["end"])

                if extend_xlim_start:
                    if self.strand == "+":
                        minn = min(minn_, minn)
                    else:
                        maxn = max(maxn_, maxn)

                if extend_xlim_end:
                    if self.strand == "+":
                        maxn = max(maxn_, maxn)
                    else:
                        minn = min(minn_, minn)

            if maxn - minn > 400:
                step = (maxn - minn) // 400 * 100  # 坐标轴步长
            else:
                step = (maxn - minn) // 40 * 10

            ax = axes[nrows - 1]
            if self.strand == "+":
                xticks = np.arange(minn, maxn + step, step)
                ax.set_xticks(xticks)
                ax.set_xticklabels(xticks - minn)
                ax.set_xlim(minn, maxn)
            else:
                xticks = np.arange(maxn, minn - step, -step)
                ax.set_xticks(xticks)
                ax.set_xticklabels(maxn - xticks)
                ax.set_xlim(minn, maxn)
                ax.invert_xaxis()

        else:
            # plot genome coordinate
            start = (self.start + 100_000) // 100_000 * 100_000
            step = (self.end - start) // 3
            step = step // 1000 * 1000
            ax = axes[nrows - 1]
            xticks = np.arange(start, self.end, step)
            ax.set_xticks(xticks)
            if start > 10_000:
                ax.set_xticklabels(map(lambda x: f'{x:,d} kb', xticks//1000))
            ax.set_xlim(self.start, self.end)

        return axes


import sys,json,requests,subprocess

def main():
	answer = requests.post("http://localhost:5279", json={"method": "get", "params": {"uri": sys.argv[1]}}).json()
	stream_type=answer["result"]["metadata"]["stream_type"]

	if stream_type != "video":
		sys.exit("Not a video file, exiting...")

	URL = answer["result"]["streaming_url"]
	print("Streaming URL:",URL )
	if "metadata" in answer["result"]:
		if "description" in answer["result"]["metadata"]:
			description=answer["result"]["metadata"]["description"]
			print("\nVideo description:\n",description, "\n")
	print("Playing with mpv...\n")
	subprocess.run(["mpv", URL])

if __name__=="__main__":
    main()


from PyQt6.QtWidgets import QApplication, QWidget, QLabel, QPushButton, QMenu, QLineEdit
import sys
from PyQt6.QtGui import QIcon, QFont, QPixmap, QMovie
from PyQt6.QtCore import QSize

app = QApplication(sys.argv)

class MyWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setFixedWidth(500)
        self.setFixedHeight(500)
        self.setWindowTitle('PersianPython')
        self.setWindowOpacity(1)
        self.setWindowIcon(QIcon('logo.png'))

        self.label = QLabel("", self)
        self.label.setGeometry(100, 200, 50, 50)

        self.button = QPushButton('QPushButton', self)
        self.button.setGeometry(200, 0, 100, 50)
        self.button.setFont(QFont('Arial', 7, QFont.Weight.ExtraBold))
        self.button.setIcon(QIcon('logo.png'))
        self.button.setIconSize(QSize(50, 50))
        self.button.clicked.connect(self.clicked)

        self.create_qline_edit()

    def create_qline_edit(self):
        self.input_text = QLineEdit(self)
        self.input_text.setGeometry(5, 5, 150, 50)
        # self.input_text.setText('Default Text !')
        self.input_text.setPlaceholderText('Enter username')
        self.input_text.setFont(QFont('Arial', 10, QFont.Weight.ExtraBold))
        self.input_text.setStyleSheet('color:blue')
        self.input_text.setEchoMode(QLineEdit.EchoMode.PasswordEchoOnEdit)

    def clicked(self):
        if self.input_text.text():
            self.label.setText(self.input_text.text())
            self.input_text.clear()

window = MyWindow()

window.show()

sys.exit(app.exec())


import cv2
import mmcv
import numpy as np
import os
import torch
import matplotlib.pyplot as plt

from mmcv import Config
from mmdet.datasets.builder import build_dataset
from mmdet.datasets.pipelines import Compose

def featuremap_2_heatmap(feature_map):
    assert isinstance(feature_map, torch.Tensor)
    feature_map = feature_map.detach()
    heatmap = feature_map[:,0,:,:]*0
    heatmaps = []
    for c in range(feature_map.shape[1]):
        heatmap+=feature_map[:,c,:,:]
    heatmap = heatmap.cpu().numpy()
    heatmap = np.mean(heatmap, axis=0)

    heatmap = np.maximum(heatmap, 0)
    heatmap /= np.max(heatmap)
    heatmaps.append(heatmap)

    return heatmaps

def retrieve_data_cfg(config_path, skip_type, cfg_options=None):
    cfg = Config.fromfile(config_path)
    if cfg_options is not None:
        cfg.merge_from_dict(cfg_options)
    # import modules from string list.
    if cfg.get('custom_imports', None):
        from mmcv.utils import import_modules_from_strings
        import_modules_from_strings(**cfg['custom_imports'])
    train_data_cfg = cfg.data.train
    train_data_cfg['pipeline'] = [
        x for x in train_data_cfg.pipeline if x['type'] not in skip_type
    ]

    return cfg

def draw_feature_map(features,save_dir = 'feature_map',name = "vis"):
    i=0
    # img = cv2.imread('data/kins/testing/image_2/000002.png')
    img_path = 'data/kins/training/image_2/000008.png'
    img_name = img_path.split("/")[-1][:-4]
    skip_type = ['DefaultFormatBundle', 'Normalize', 'Collect', 'RandomFlip','LoadAnnotations']
    config_kins = 'configs/kins/atss_r50_fpn_2x_kins.py'
    kins_cfg = retrieve_data_cfg(config_kins, skip_type)
    test_pip = Compose(kins_cfg.data.train.pipeline)
    # data = dict(img=img)
    data = dict(img_info=dict(filename=img_path), img_prefix=None)
    img =test_pip(data)['img']

    if isinstance(features,torch.Tensor):
        for heat_maps in features:
            heat_maps=heat_maps.unsqueeze(0)
            heatmaps = featuremap_2_heatmap(heat_maps)
            # 这里的h,w指的是你想要把特征图resize成多大的尺寸
            # heatmap = cv2.resize(heatmap, (h, w))
            for heatmap in heatmaps:
                heatmap = np.uint8(255 * heatmap)
                # 下面这行将热力图转换为RGB格式 ，如果注释掉就是灰度图
                heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
                superimposed_img = heatmap
                plt.imshow(superimposed_img,cmap='gray')
                plt.show()
    else:
        for featuremap in features:
            heatmaps = featuremap_2_heatmap(featuremap)
            for heatmap in heatmaps:
                heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))  # 将热力图的大小调整为与原始图像相同
                heatmap = np.uint8(255 * (1-heatmap))  # 将热力图转换为RGB格式
                heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
                superimposed_img = heatmap * 0.5 + img*0.3
                # superimposed_img = heatmap
                # plt.imshow(superimposed_img)
                # plt.imshow(superimposed_img,cmap='plasma_r')
                # plt.show()
                # 下面这些是对特征图进行保存，使用时取消注释
                # cv2.imshow("1",superimposed_img)
                # cv2.waitKey(0)
                # cv2.destroyAllWindows()
                cv2.imwrite(os.path.join(save_dir,name + "_" + img_name + "_" + str(i) +'.png'), superimposed_img)
                i=i+1


# -*- coding: utf-8 -*-
"""
Created on %(date)s

@author: %(username)s
"""
import os
import timeit
import time

import numpy as np

import pyximport
pyximport.install()
from .idw_nebs import slct_nebrs_cy, get_idw

from krig_single_pts import (OrdinaryKriging,
                             SimpleKriging,
                             ExternalDriftKriging_MD)

def get_all_grid(interp_x_coords_arr,
                 interp_y_coords_arr,
                 nebor_x_coords_arr,
                 nebor_y_coords_arr,
                 nebor_vals_arr,
                 idw_exp,
                 n_quads,
                 n_per_quad,
                 min_dist_thresh,
                 n_nebs,
                 interp_types,
                 vg_model=None,
                 use_close_nebs=True):

    n_interp_pts = interp_x_coords_arr.shape[0]
    n_interp_types = len(interp_types)
    interp_vals_arr = np.full((n_interp_pts, n_interp_types), np.nan)

    if 'idw' in interp_types:
        idw_idx = interp_types.index('idw')

    if 'ok' in interp_types:
        assert vg_model
        ok_idx = interp_types.index('ok')

#    if 'sk' in in_interp_types:
#        assert vg_model
#        sk_idx = interp_types.index('sk')
#
#    if 'edk' in in_interp_types:
#        assert vg_model
#        edk_idx = interp_types.index('edk')

    if vg_model is not None:
        vg_model = str(vg_model)

    if use_close_nebs:
        # arrays to hold temporary values
        dists_arr = np.zeros(n_nebs, dtype=np.float64)
        prcssed_nebrs_arr = np.zeros(n_nebs, dtype=np.uint32)
        slctd_nebrs_arr = np.zeros(n_nebs, dtype=np.uint32)
        slctd_nebrs_dists_arr = np.zeros(n_nebs, dtype=np.float64)
        nebs_idxs_arr = np.zeros(n_nebs, dtype=np.int64)
        fin_nebs_idxs_arr = np.zeros(n_nebs, dtype=np.uint32)

    for i in range(n_interp_pts):
        x = interp_x_coords_arr[i]
        y = interp_y_coords_arr[i]

        if use_close_nebs:
            slct_nebrs_cy(x,
                          y,
                          nebor_x_coords_arr,
                          nebor_y_coords_arr,
                          n_quads,
                          n_per_quad,
                          min_dist_thresh,
                          n_nebs,
                          prcssed_nebrs_arr,
                          slctd_nebrs_arr,
                          nebs_idxs_arr,
                          dists_arr,
                          slctd_nebrs_dists_arr,
                          fin_nebs_idxs_arr)

            fin_nebs_idxs_bool_arr = fin_nebs_idxs_arr.astype(bool)
            fin_nebrs_x_crds = nebor_x_coords_arr[fin_nebs_idxs_bool_arr]
            fin_nebrs_y_crds = nebor_y_coords_arr[fin_nebs_idxs_bool_arr]
            fin_nebrs_vals = nebor_vals_arr[fin_nebs_idxs_bool_arr]

        else:
            fin_nebrs_x_crds = nebor_x_coords_arr
            fin_nebrs_y_crds = nebor_y_coords_arr
            fin_nebrs_vals = nebor_vals_arr

        if fin_nebrs_x_crds.shape[0]:
            if 'idw' in interp_types:
                idw_val = get_idw(x,
                                  y,
                                  fin_nebrs_x_crds,
                                  fin_nebrs_y_crds,
                                  fin_nebrs_vals,
                                  idw_exp)
                interp_vals_arr[i, idw_idx] = idw_val

            if 'ok' in interp_types:
                _x = np.array([x])
                _y = np.array([y])

#                import IPython.core.debugger
#                dbg = IPython.core.debugger.Pdb()
#                dbg.set_trace()

                ord_krig = OrdinaryKriging(xi=fin_nebrs_x_crds,
                                           yi=fin_nebrs_y_crds,
                                           zi=fin_nebrs_vals,
                                           xk=_x,
                                           yk=_y,
                                           model=vg_model)

                ord_krig.krige()

                interp_vals_arr[i, ok_idx] = ord_krig.zk[0]

        else:
            print('No nebors!')

    return interp_vals_arr

if __name__ == '__main__':
    import pandas as pd
    import matplotlib.pyplot as plt

    plt.ioff()

#    print('\a\a\a\a Started on %s \a\a\a\a\n' % time.asctime())
#    START = timeit.default_timer()  # to get the runtime of the program

    main_dir = os.path.join(r'P:\Synchronize\Python3Codes',
                            r'krig_idw_nebs\test_ppt')

    in_data_file = r'infilled_var_df_infill_stns.csv'
    in_coords_file = r'infilled_var_df_infill_stns_coords.csv'

    in_vgs_file = (r'ppt_fitted_variograms__nk_1.000__evg_name_'
                   r'robust__ngp_5__h_typ_var.csv')

    sep = ';'
    time_fmt = '%Y-%m-%d'

    interp_date = '2014-06-05'

    interp_types = ['idw', 'ok']#, 'sk', 'edk']

    idw_exp = 3
    cell_size = 5000
    buff_cells = 2
#    n_nebs = 10
    n_quads = 20
    n_per_quad = 1
    min_dist_thresh = 100.0

    os.chdir(main_dir)

    in_data_df = pd.read_csv(in_data_file, sep=sep, index_col=0)
    in_coords_df = pd.read_csv(in_coords_file, sep=sep, index_col=0)
    in_vgs_df = pd.read_csv(in_vgs_file, sep=sep, index_col=0)

    in_data_df.index = pd.to_datetime(in_data_df.index, format=time_fmt)
    in_coords_df = in_coords_df[['X', 'Y']]
    in_vgs_df = in_vgs_df.iloc[:, 0]

    curr_vals_ser = in_data_df.loc[interp_date].dropna().iloc[:100].copy()
    curr_coords_df = in_coords_df.loc[curr_vals_ser.index].dropna().copy()
    vg_model = in_vgs_df.loc[interp_date]

    comm_stns = curr_coords_df.index
    comm_stns = comm_stns.intersection(curr_vals_ser.index)

    curr_vals_ser = curr_vals_ser.loc[comm_stns]
    curr_coords_df = curr_coords_df.loc[comm_stns]

    curr_x_vals = curr_coords_df['X'].values.astype(np.float64, order='C')
    curr_y_vals = curr_coords_df['Y'].values.astype(np.float64, order='C')
    curr_vals = curr_vals_ser.values.astype(np.float64, order='C')

    assert np.all(curr_x_vals)
    assert np.all(curr_y_vals)
    assert np.all(np.logical_not(np.isnan(curr_vals)))

    min_x = curr_x_vals.min()
    max_x = curr_x_vals.max()

    min_y = curr_y_vals.min()
    max_y = curr_y_vals.max()

    min_val = curr_vals.min()
    max_val = curr_vals.max()

    n_nebs = curr_vals.shape[0]

    x_interp_vec = np.arange(min_x - (buff_cells * cell_size),
                             max_x + (buff_cells * cell_size), cell_size)

    y_interp_vec = np.arange(min_y - (buff_cells * cell_size),
                             max_y + (buff_cells * cell_size), cell_size)

    interp_x_coords, interp_y_coords = np.meshgrid(x_interp_vec, y_interp_vec)

    orig_grid_shape = interp_x_coords.shape

    interp_x_coords_arr = interp_x_coords.ravel()
    interp_y_coords_arr = interp_y_coords.ravel()

    pcol_x_coords = np.linspace(x_interp_vec[0],
                                x_interp_vec[-1],
                                x_interp_vec.shape[0] + 1)

    pcol_y_coords = np.linspace(y_interp_vec[0],
                                y_interp_vec[-1],
                                y_interp_vec.shape[0] + 1)

    (pcol_x_mesh_coords,
     pcol_y_mesh_coords) = np.meshgrid(pcol_x_coords, pcol_y_coords)

    atleast_one = 0

    if 'idw' in interp_types:
        idw_idx = interp_types.index('idw')
        atleast_one += 1

    if 'ok' in interp_types:
        assert vg_model
        ok_idx = interp_types.index('ok')
        atleast_one += 1

    if 'sk' in interp_types:
        assert vg_model
        sk_idx = interp_types.index('sk')
        atleast_one += 1

    if 'edk' in interp_types:
        assert vg_model
        edk_idx = interp_types.index('edk')
        atleast_one += 1

    assert atleast_one

    print('\a\a\a\a Started on %s \a\a\a\a\n' % time.asctime())
    START = timeit.default_timer()  # to get the runtime of the program

    with_neb_sel_interps = get_all_grid(interp_x_coords_arr,
                                        interp_y_coords_arr,
                                        curr_x_vals,
                                        curr_y_vals,
                                        curr_vals,
                                        idw_exp,
                                        n_quads,
                                        n_per_quad,
                                        min_dist_thresh,
                                        n_nebs,
                                        interp_types,
                                        vg_model,
                                        use_close_nebs=True)

    wo_neb_sel_interps = get_all_grid(interp_x_coords_arr,
                                      interp_y_coords_arr,
                                      curr_x_vals,
                                      curr_y_vals,
                                      curr_vals,
                                      idw_exp,
                                      n_quads,
                                      n_per_quad,
                                      min_dist_thresh,
                                      n_nebs,
                                      interp_types,
                                      vg_model,
                                      use_close_nebs=False)

    fig, axes = plt.subplots(atleast_one, atleast_one)

    if 'idw' in interp_types:
        idw_neb_sel_interp_grid = \
            with_neb_sel_interps[:, idw_idx].reshape(orig_grid_shape)
        idw_wo_neb_sel_interp_grid = \
            wo_neb_sel_interps[:, idw_idx].reshape(orig_grid_shape)

        axes[idw_idx, 0].pcolormesh(pcol_x_mesh_coords,
                                    pcol_y_mesh_coords,
                                    idw_neb_sel_interp_grid,
                                    vmin=min_val,
                                    vmax=max_val)
        axes[idw_idx, 0].scatter(curr_x_vals, curr_y_vals)

        axes[idw_idx, 1].pcolormesh(pcol_x_mesh_coords,
                                    pcol_y_mesh_coords,
                                    idw_wo_neb_sel_interp_grid,
                                    vmin=min_val,
                                    vmax=max_val)
        axes[idw_idx, 1].scatter(curr_x_vals, curr_y_vals)

    if 'ok' in interp_types:
        ok_neb_sel_interp_grid = \
            with_neb_sel_interps[:, ok_idx].reshape(orig_grid_shape)
        ok_wo_neb_sel_interp_grid = \
            wo_neb_sel_interps[:, ok_idx].reshape(orig_grid_shape)

        axes[ok_idx, 0].pcolormesh(pcol_x_mesh_coords,
                                   pcol_y_mesh_coords,
                                   ok_neb_sel_interp_grid,
                                   vmin=min_val,
                                   vmax=max_val)
        axes[ok_idx, 0].scatter(curr_x_vals, curr_y_vals)

        axes[ok_idx, 1].pcolormesh(pcol_x_mesh_coords,
                                   pcol_y_mesh_coords,
                                   ok_wo_neb_sel_interp_grid,
                                   vmin=min_val,
                                   vmax=max_val)
        axes[ok_idx, 1].scatter(curr_x_vals, curr_y_vals)

    plt.show()
    STOP = timeit.default_timer()  # Ending time
    print(('\n\a\a\a Done with everything on %s. Total run time was'
           ' about %0.4f seconds \a\a\a' % (time.asctime(), STOP-START)))


# -*- encoding: utf-8 -*-
"""
@Author  : QADai
@Email   : daibiao83@126.com
"""
import pytest
import sys
from datetime import datetime
sys.path.append(r'../../')
from configure.configure import *
from src.base.browser import *
from src.base.browser_engine import *

@pytest.mark.hookwrapper
def pytest_runtest_makereport(item):
    """
    Extends the PyTest Plugin to take and embed screenshot in html report, whenever test fails.
    :param item:
    """
    pytest_html = item.config.pluginmanager.getplugin('html')
    outcome = yield
    report = outcome.get_result()
    extra = getattr(report, 'extra', [])

    if report.when == 'call' or report.when == "setup":
        xfail = hasattr(report, 'wasxfail')
        if (report.skipped and xfail) or (report.failed and not xfail):
            dirpng = screen_path
            if os.path.exists(dirpng) and os.path.isdir(dirpng):
                pass
            else:
                os.mkdir(dirpng)
            i = datetime.now()
            screen_name = str(i.year) + str(i.month) + str(i.day) + str(i.hour) + str(i.minute) + str(i.second)

            file_name = dirpng + report.nodeid.replace("::", "_") + screen_name +".png"
            file_name1 = r'./png/' + report.nodeid.replace("::", "_") + screen_name + ".png"
            _capture_screenshot(file_name)
            if file_name:
                html = '

' % file_name1
                extra.append(pytest_html.extras.html(html))
        report.extra = extra
        report.nodeid = report.nodeid.encode("utf-8").decode("unicode_escape")

def _capture_screenshot(name):
    driver = Webdriver_Opts(current_using_browser)
    browser = SeleniumDriver(driver.get_webdriver_instance())
    browser.capture_screen(name)

@pytest.fixture(scope='session', autouse=True)
def browser():
    global driver
    if driver is None:
        driver = Webdriver_Opts(current_using_browser)
    return driver


from base.base_model import BaseModel
from keras.layers import Dense,GlobalAveragePooling2D
from keras.applications import mobilenet

class MobileNet_model(BaseModel):
    def __init__(self, config):
        super(MobileNet_model, self).__init__(config)
        self.build_model()

    def build_model(self):

        self.model = mobilenet.MobileNet(input_shape=None, alpha=1.0, depth_multiplier=1, dropout=1e-3, include_top=True, weights='imagenet', input_tensor=None, pooling=None, classes=1000)

        #  This compiles the model architecture and the necessary functions that we
        #  categorical crossentropy is the loss function for classification problems with more than 2 classes
        self.model.compile(loss='categorical_crossentropy',
                      optimizer=self.config.model.optimizer,
                      metrics=['accuracy']) # want to see how accurate the model is (but are minimize the loss)


# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import numpy as np
import matplotlib.pyplot as mp
outcomes = np.random.hypergeometric(25, 1, 3, 100)
scores = [0]
for outcome in outcomes:
    if outcome == 3:
        scores.append(scores[-1] + 1)
    else:
        scores.append(scores[-1] - 6)
scores = np.array(scores)
mp.figure('Hypergeometric',
          facecolor='lightgray')
mp.title('Hypergeometric', fontsize=20)
mp.xlabel('Round', fontsize=14)
mp.ylabel('Score', fontsize=14)
mp.tick_params(labelsize=10)
mp.grid(linestyle=':')
if scores[-1] < scores[0]:
    color = 'limegreen'
elif scores[-1] > scores[0]:
    color = 'orangered'
else:
    color = 'dodgerblue'
mp.plot(scores, c=color, label='Score')
mp.legend()
mp.show()


from PIL import Image, ImageDraw, ImageFont
import concurrent.futures
import random
import os
import time
from math import ceil

# this file will create the necessary images, and then return that image based on a list of texts
# has all the helper methods

# creating all fonts
all_fonts = [f'./yearbook/fonts/{f}' for f in os.listdir('./yearbook/fonts/')]

# creates the base 64 version of images
def create_base64_images(image):
    import base64
    from io import BytesIO

    buffered = BytesIO()
    image.save(buffered, format="PNG")
    encoded = base64.b64encode(buffered.getvalue())

    encoded = encoded.decode("utf-8")

    mime = "image/jpeg"
    uri = "data:%s;base64, %s" % (mime, encoded)
    return uri

# texts is a list that contains both messages and authors
def create_images(texts=[], size=(1044, 1044)):

    # now that we know we can resize stuff, we want to resize it so that we can fit 4 on one page
    txt = Image.new('RGBA', size, (255, 255, 255, 255))
    fnt = ImageFont.truetype(all_fonts[0], 40)
    d = ImageDraw.Draw(txt)

    # we need a method to just grab all the processed texts
    messages = []
    for text in texts:
        messages.append(process_images((text, txt, fnt, d, size)))

    messages = list(filter(lambda msg: msg, messages))

    images_out = []

    # if messages exist, do img processing on them
    if messages:

        for j in range(ceil(len(messages)/4)):

            current_time = time.time()  # get the current time

            # counts the number of actual messages
            count = 0

            # separate the messages into groups of four
            msgs = []

            # loops through, getting all the messages
            for k in range(4):
                if len(messages) > k + j * 4:
                    msgs.append(messages[k+j*4][0])
                    count += 1
                else:
                    msgs.append("")

            # generates the sizes of all the texts
            sizes_of_texts = [d.multiline_textsize(msgs[i], font=fnt)[
                1] for i in range(4)]

            # creating the right sizes now
            new_size = (1044,
                        max(sizes_of_texts[0] + sizes_of_texts[2],
                            sizes_of_texts[1] + sizes_of_texts[3]) + 60)

            # creates new images with new sizes
            txt = Image.new('RGBA', new_size, (255, 255, 255, 255))
            d = ImageDraw.Draw(txt)

            for i in range(count):  # loops through all the

                fnt = ImageFont.truetype(all_fonts[random.randint(0, len(all_fonts)-1)], 40)

                # sets the right positioning
                pos_x = 10 if (i % 2) == 0 else new_size[0]/2
                pos_y = 10 if int(
                    (i % 4) / 2) == 0 else new_size[1] - sizes_of_texts[i % 2 + 1] - 10

                # better generation of colours, so that they aren't too light - their sum has to be < 500
                red = random.randint(0, 255)
                green = random.randint(0, 500 - red)
                blue = random.randint(0, 500 - red - green)

                # writes in the correct text
                d.multiline_text((pos_x, pos_y), msgs[i], font=fnt, fill=(
                    red, green, blue, 255))

            images_out.append(create_base64_images(txt))

    return images_out

# do processing of images here now, here is the method to call:

def process_images(messages=()):

    whole_msg, txt, fnt, d, size = messages  # gets the message and author here
    msg, author = whole_msg

    # gets the message
    message = find_best_string(message=msg.split(' '), pos=0, font=fnt, d=d, size_x=(
        size[0]-20)/2, size_y=(size[1]-20)/2, memo={})

    if message == 'Could not resize it!':
        return None

    return (message, author)

# Steps for DP:
# 1. Define Subproblem, just a suffix [i:]
# 2. Guess, where to place a space
# 3. Relate Subproblems, min(area of one, vs area of another)
# 4. Memoize or Tabulate
# 5. Return Orig. Problem

def find_best_string(message=[], pos=0, font=None, d=None, size_x=300, size_y=300, memo={}):

    # if it is already memoized, it goes here
    if memo.get(pos):
        return memo.get(pos)

    # it's -1 and not len(message) because we don't care if the last one has a new line, 1/2 the time
    if pos == len(message) - 1:
        memo[pos] = message[pos]
        return message[pos]

    # creates both possibilities
    answers = [f'''{message[pos]}\n''' + find_best_string(message=message,
                                                          pos=pos+1, font=font, d=d, memo=memo),
               f'''{message[pos]} ''' + find_best_string(message=message,
                                                         pos=pos+1, font=font, d=d, memo=memo)]

    # if it is not resizeable for one of them, returns the other one right away
    if answers[0] == 'Could not resize it!':
        memo[pos] = answers[1]
        return answers[1]

    elif answers[1] == 'Could not resize it!':
        memo[pos] = answers[0]
        return answers[0]

    # finds the area of them, to check if it can resize
    area0 = d.multiline_textsize(answers[0], font=font)
    area1 = d.multiline_textsize(answers[1], font=font)

    # shows that it can't be resized here, if one can't, then just returns the other
    if (area0[0] > size_x and area1[0] > size_x) or (area0[1] > size_y and area1[1] > size_y):
        memo[pos] = 'Could not resize it!'
        return 'Could not resize it!'
    elif area0[0] > size_x or area0[1] > size_y:
        memo[pos] = answers[1]
        return answers[1]
    elif area1[0] > size_x or area1[1] > size_y:
        memo[pos] = answers[0]
        return answers[0]

    # the optimal choice is always the one with the least number of lines
    memo[pos] = answers[1]
    return answers[1]


import time, math
from selenium import webdriver as wd
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By

link = 'http://suninjuly.github.io/execute_script.html'

try:
    service = Service(executable_path='C:\chromedriver\chromedriver.exe')
    browser = wd.Chrome()
    browser.get(link)

    x = browser.find_element(By.XPATH, "//span[contains(@id, 'input_value')]").text

    browser.find_element(By.ID, 'answer').send_keys(str(math.log(abs(12*math.sin(int(x))))))

    browser.find_element(By.CSS_SELECTOR, '#robotCheckbox').click()

    robotsRule = browser.find_element(By.XPATH, '//*[@id="robotsRule"]')
    button = browser.find_element(By.XPATH, "//button[contains(@type, 'submit')]")
    browser.execute_script("return arguments[0].scrollIntoView(true)", robotsRule)
    robotsRule.click(); time.sleep(1); button.click()

except Exception as ex: print(ex)
finally: time.sleep(15); browser.quit()


from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals

import logging.config
import os
import shutil
import sys
from sqlite3 import DatabaseError as SQLite3DatabaseError

import django
from diskcache.fanout import FanoutCache
from django.conf import settings
from django.core.exceptions import ValidationError
from django.core.management import call_command
from django.core.management.base import handle_default_options
from django.db.utils import DatabaseError

import kolibri
from kolibri.core.device.utils import device_provisioned
from kolibri.core.device.utils import provision_from_file
from kolibri.core.deviceadmin.exceptions import IncompatibleDatabase
from kolibri.core.deviceadmin.utils import get_backup_files
from kolibri.core.upgrade import matches_version
from kolibri.core.upgrade import run_upgrades
from kolibri.core.utils.cache import process_cache
from kolibri.deployment.default.sqlite_db_names import ADDITIONAL_SQLITE_DATABASES
from kolibri.plugins.utils import autoremove_unavailable_plugins
from kolibri.plugins.utils import check_plugin_config_file_location
from kolibri.plugins.utils import enable_new_default_plugins
from kolibri.plugins.utils import run_plugin_updates
from kolibri.utils.conf import KOLIBRI_HOME
from kolibri.utils.conf import LOG_ROOT
from kolibri.utils.conf import OPTIONS
from kolibri.utils.debian_check import check_debian_user
from kolibri.utils.logger import get_base_logging_config
from kolibri.utils.sanity_checks import check_content_directory_exists_and_writable
from kolibri.utils.sanity_checks import check_database_is_migrated
from kolibri.utils.sanity_checks import check_default_options_exist
from kolibri.utils.sanity_checks import check_django_stack_ready
from kolibri.utils.sanity_checks import check_log_file_location
from kolibri.utils.sanity_checks import DatabaseInaccessible
from kolibri.utils.sanity_checks import DatabaseNotMigrated
from kolibri.utils.sanity_checks import ensure_job_tables_created
from kolibri.utils.server import get_status
from kolibri.utils.server import NotRunning

logger = logging.getLogger(__name__)

def version_file():
    """
    During test runtime, this path may differ because KOLIBRI_HOME is
    regenerated
    """
    return os.path.join(KOLIBRI_HOME, ".data_version")

def version_updated(kolibri_version, version_file_contents):
    return kolibri_version != version_file_contents

def should_back_up(kolibri_version, version_file_contents):
    change_version = kolibri_version != version_file_contents
    return (
        # Only back up if there was a previous version
        version_file_contents
        # That version has changed
        and change_version
        # The previous version was not a dev version
        and "dev" not in version_file_contents
        # And the new version is not a dev version
        and "dev" not in kolibri_version
    )

def conditional_backup(kolibri_version, version_file_contents):
    if should_back_up(kolibri_version, version_file_contents):
        # Non-dev version change, make a backup no matter what.
        from kolibri.core.deviceadmin.utils import dbbackup
        from kolibri.core.deviceadmin.utils import default_backup_folder

        try:
            backup = dbbackup(version_file_contents)
            default_path = default_backup_folder()
            backups = get_backup_files()
            if len(backups) > 2:
                for old_backup in backups[2:]:
                    os.remove(os.path.join(default_path, old_backup))
            logger.info("Backed up database to: {path}".format(path=backup))
        except IncompatibleDatabase:
            logger.warning(
                "Skipped automatic database backup, not compatible with "
                "this DB engine."
            )

def get_version():
    try:
        version = open(version_file(), "r").read()
        return version.strip() if version else ""
    except IOError:
        return ""

def _migrate_databases():
    """
    Try to migrate all active databases. This should not be called unless Django has
    been initialized.
    """
    for database in settings.DATABASES:
        call_command("migrate", interactive=False, database=database)

    # load morango fixtures needed for certificate related operations
    call_command("loaddata", "scopedefinitions")

class DefaultDjangoOptions(object):
    __slots__ = ["settings", "pythonpath"]

    def __init__(self, settings, pythonpath):
        self.settings = settings
        self.pythonpath = pythonpath

def setup_logging(debug=False, debug_database=False):
    """
    Configures logging in cases where a Django environment is not supposed
    to be configured.
    """
    # Sets the global DEBUG flag to be picked up in other contexts
    # (Django settings)
    os.environ["KOLIBRI_DEBUG"] = str(debug)
    os.environ["KOLIBRI_DEBUG_LOG_DATABASE"] = str(debug_database)

    # Would be ideal to use the upgrade logic for this, but that is currently
    # only designed for post-Django initialization tasks. If there are more cases
    # for pre-django initialization upgrade tasks, we can generalize the logic here
    if matches_version(get_version(), "<0.12.4"):
        check_log_file_location()

    LOGGING = get_base_logging_config(
        LOG_ROOT, debug=debug, debug_database=debug_database
    )
    logging.config.dictConfig(LOGGING)

def _setup_django():
    """
    Do our django setup - separated from initialize to reduce complexity.
    """
    try:
        django.setup()

    except (DatabaseError, SQLite3DatabaseError) as e:
        if "malformed" in str(e):
            logger.error(
                "Your database appears to be corrupted. If you encounter this,"
                "please immediately back up all files in the .kolibri folder that"
                "end in .sqlite3, .sqlite3-shm, .sqlite3-wal, or .log and then"
                "contact Learning Equality. Thank you!"
            )
        raise

def _copy_preseeded_db(db_name, target=None):
    target = target or "{}.sqlite3".format(db_name)
    target = os.path.join(KOLIBRI_HOME, target)
    if not os.path.exists(target):
        try:
            import kolibri.dist

            db_path = os.path.abspath(
                os.path.join(
                    os.path.dirname(kolibri.dist.__file__),
                    "home/{}.sqlite3".format(db_name),
                )
            )
            shutil.copyfile(db_path, target)
            logger.info(
                "Copied preseeded database from {} to {}".format(db_path, target)
            )
        except (ImportError, IOError, OSError):
            logger.warning(
                "Unable to copy pre-migrated database from {} to {}".format(
                    db_path, target
                )
            )

def _upgrades_before_django_setup(updated, version):
    if version and updated:
        check_plugin_config_file_location(version)

    # Do this here so that we can fix any issues with our configuration file before
    # we attempt to set up django.
    autoremove_unavailable_plugins()

    if updated:
        # Reset the enabled plugins to the defaults
        # This needs to be run before dbbackup because
        # dbbackup relies on settings.INSTALLED_APPS
        enable_new_default_plugins()

    # Ensure that we have done all manipulations of our plugins registry before
    # we do the check for options.ini as that will invoke our plugin registry.
    # Check if there is an options.ini file exist inside the KOLIBRI_HOME folder
    check_default_options_exist()

    if OPTIONS["Database"]["DATABASE_ENGINE"] == "sqlite":
        # If we are using sqlite,
        # we can shortcut migrations by using the preseeded databases
        # that we bundle in the Kolibri whl file.
        if not version:
            logger.info("Attempting to setup using pre-migrated databases")
            # Only copy the default database if this is a fresh install
            _copy_preseeded_db("db", target=OPTIONS["Database"]["DATABASE_NAME"])

        if not version or updated:
            # If this is an upgrade, it is possible we've added an additional
            # database, so we can attempt to copy a preseeded database here.
            for db_name in ADDITIONAL_SQLITE_DATABASES:
                _copy_preseeded_db(db_name)

def _post_django_initialization():
    # Import here to prevent the module level access to Kolibri options
    # which causes premature registration of Kolibri plugins.
    from kolibri.deployment.default.cache import CACHES

    if "process_cache" in CACHES:  # usually it means not using redis
        if "DatabaseCache" not in CACHES["process_cache"]["BACKEND"]:
            try:
                process_cache.cull()
            except SQLite3DatabaseError:
                shutil.rmtree(process_cache.directory, ignore_errors=True)
                os.mkdir(process_cache.directory)
                process_cache._cache = FanoutCache(
                    process_cache.directory,
                    settings.CACHES["process_cache"]["SHARDS"],
                    settings.CACHES["process_cache"]["TIMEOUT"],
                    **settings.CACHES["process_cache"]["OPTIONS"]
                )

def _upgrades_after_django_setup(updated, version):
    # If device is not provisioned, attempt automatic provisioning
    if not device_provisioned() and OPTIONS["Paths"]["AUTOMATIC_PROVISION_FILE"]:
        try:
            provision_from_file(OPTIONS["Paths"]["AUTOMATIC_PROVISION_FILE"])
        except ValidationError as e:
            logging.error(
                "Tried to automatically provision the device but received an error"
            )
            logging.error(e)

def initialize(  # noqa C901
    skip_update=False,
    settings=None,
    debug=False,
    debug_database=False,
    no_input=True,
    pythonpath=None,
):
    """
    This should be called before starting the Kolibri app, it initializes Kolibri plugins
    and sets up Django.
    """
    check_debian_user(no_input)

    setup_logging(debug=debug, debug_database=debug_database)

    default_options = DefaultDjangoOptions(settings, pythonpath)

    handle_default_options(default_options)

    version = get_version()

    updated = version_updated(kolibri.__version__, version)

    if not skip_update:
        _upgrades_before_django_setup(updated, version)

        try:
            ensure_job_tables_created()
        except Exception as e:
            logging.error(
                "The job tables were not fully migrated. Tried to "
                "create them in the database and an error occurred: "
                "{}".format(e)
            )
            raise

    _setup_django()

    _post_django_initialization()

    if updated and not skip_update:
        conditional_backup(kolibri.__version__, version)

        if version:
            logger.info(
                "Version was {old}, new version: {new}".format(
                    old=version, new=kolibri.__version__
                )
            )
        else:
            logger.info("New install, version: {new}".format(new=kolibri.__version__))
        update(version, kolibri.__version__)

    check_content_directory_exists_and_writable()

    if not skip_update:
        # Run any plugin specific updates here in case they were missed by
        # our Kolibri version based update logic.
        run_plugin_updates()

        check_django_stack_ready()

        try:
            check_database_is_migrated()
        except DatabaseNotMigrated:
            try:
                _migrate_databases()
            except Exception as e:
                logging.error(
                    "The database was not fully migrated. Tried to "
                    "migrate the database and an error occurred: "
                    "{}".format(e)
                )
                raise
        except DatabaseInaccessible as e:
            logging.error(
                "Tried to check that the database was accessible "
                "and an error occurred: {}".format(e)
            )
            raise

        _upgrades_after_django_setup(updated, version)

def update(old_version, new_version):
    """
    Called whenever a version change in kolibri is detected
    """

    logger.info("Running update routines for new version...")

    try:
        # Check if there are other kolibri instances running
        # If there are, then we need to stop users from starting kolibri again.
        get_status()
        logger.error(
            "There is a Kolibri server running. "
            "Running updates now could cause a database error. "
            "Please use `kolibri stop` and try again. "
        )
        sys.exit(1)

    except NotRunning:
        pass

    _migrate_databases()

    run_upgrades(old_version, new_version)

    with open(version_file(), "w") as f:
        f.write(kolibri.__version__)


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import plot
from matplotlib.animation import FuncAnimation

from params import QUEST_TIMES, CRAFT_TIMES, PR_DROP_LOOT_FROM_QUEST
from legion import Legion

def recruit_drop_rates(
    N = 100,
    k = 100,
    b = 1,
    d = 0,
    f = lambda N,k: N/k
):
    """
        N: recruit questing population size
        b: birth rate
        d: death rate
        k: carrying capacity (crafting population)
    """
    return (b * f(N, k) - d) * N

def linear(N=100, k=100):
    return (1 - N/k)

def inverse_linear(N=100, k=100):
    return 1 / (1 + N/k)

def inverse_quad(N=100, k=100, s=1, D=3):
    ONE = 100_000
    return ONE / (ONE + ((N*ONE)/(k*ONE*s))**2*ONE) / D

def inverse_exp(N=100, k=100):
    return np.exp(-np.log(2) * N / k)

ax1 = 1
ax2 = 1
FRAMES = 100

def init_plot(i=0):
    # do nothing, prevents FuncAnim calling initialization twice
    return

def draw_recruit_drop_paths(i):

    num_craftor = 8 + i * 8 # each i-frame is 8 craftors
    num_recruits_questing = 16 + i * 16 # each i-frame is 8 craftors

    # clear plots to redraw
    ax2.clear()
    ax1.clear()
    # ax2.clear()

    colors = [
        'crimson',
        'mediumorchid',
        'royalblue',
        'black',
    ]

    N = np.linspace(0, 2000, 2001)
    N_k200 = [n/200 for n in N]
    N_k400 = [n/400 for n in N]
    N_k800 = [n/800 for n in N]

    growth_rate_1 = [recruit_drop_rates(N=n, k=200, f=inverse_quad) for n in N]
    growth_rate_2 = [recruit_drop_rates(N=n, k=400, f=inverse_quad) for n in N]
    growth_rate_3 = [recruit_drop_rates(N=n, k=800, f=inverse_quad) for n in N]

    success_rate_1 = [inverse_quad(N=n, k=200) for n in N]
    success_rate_2 = [inverse_quad(N=n, k=400) for n in N]
    success_rate_3 = [inverse_quad(N=n, k=800) for n in N]

    # array index for the growth_rate of current num_recruits_questing
    index_growth_num_recruits_questing = np.argmin(np.abs(np.subtract(N, num_recruits_questing)))

    treasure_growth_yvalue = {
        '1': growth_rate_1[index_growth_num_recruits_questing],
        '2': growth_rate_2[index_growth_num_recruits_questing],
        '3': growth_rate_3[index_growth_num_recruits_questing],
    }
    drop_rate_yvalue = {
        '1': success_rate_1[index_growth_num_recruits_questing],
        '2': success_rate_2[index_growth_num_recruits_questing],
        '3': success_rate_3[index_growth_num_recruits_questing],
    }

    ############## PLOTS ##################

    ####### change in summoning success rate, vs population
    ax1.plot(N, success_rate_1 , color=colors[0])
    ax1.plot(N, success_rate_2 , color=colors[1])
    ax1.plot(N, success_rate_3 , color=colors[2])

    ax2.set(xlabel='#recruits', ylabel='Number of treasure fragments created (~7days) ')
    ax1.set(xlabel='#recruits', ylabel='Fragment drop rate')

    ax1.plot([num_recruits_questing], [drop_rate_yvalue['1']],
        label=r"Pr(recruits|craftors=200): %{:.1%}".format(drop_rate_yvalue['1']),
        color=colors[0], linestyle="-", marker="*")

    ax1.plot([num_recruits_questing], [drop_rate_yvalue['2']],
        label=r"Pr(recruits|craftors=400): %{:.1%}".format(drop_rate_yvalue['2']),
        color=colors[1], linestyle="-", marker="*")

    ax1.plot([num_recruits_questing], [drop_rate_yvalue['3']],
        label=r"Pr(recruits|craftors=800): %{:.1%}".format(drop_rate_yvalue['3']),
        color=colors[2], linestyle="-", marker="*")

    ####### change in growth rate, vs population
    ax2.plot(N, growth_rate_1, color=colors[0])
    ax2.plot([num_recruits_questing], [treasure_growth_yvalue['1']],
        label="200 Craftors => {:.0f} Fragments dropped/period".format(treasure_growth_yvalue['1']),
        color=colors[0], marker="*")

    ax2.plot(N, growth_rate_2, color=colors[1])
    ax2.plot([num_recruits_questing], [treasure_growth_yvalue['2']],
        label="400 Craftors => {:.0f} Fragments dropped/period".format(treasure_growth_yvalue['2']),
        color=colors[1], marker="*")

    ax2.plot(N, growth_rate_3, color=colors[2])
    ax2.plot([num_recruits_questing], [treasure_growth_yvalue['3']],
        label="800 Craftors => {:.0f} Fragments dropped/period".format(treasure_growth_yvalue['3']),
        color=colors[2], marker="*")

    ########################

    ax2.axvline(x=num_recruits_questing, color='black', linestyle=':', alpha=0.5)
    ax1.axvline(x=num_recruits_questing, color='black', linestyle=':', alpha=0.5)

    ax2.set_title('Recruits finding fragments per Period (~7days) | {} Recruits'.format(num_recruits_questing), size=10)
    ax1.set_title('% Fragment droprate | {} Recruits'.format(num_recruits_questing), size=10)

    ax2.legend()
    ax1.legend()

    ax1.legend(bbox_to_anchor=(1.42, 1), loc="upper right")
    ax2.legend(bbox_to_anchor=(1.48, 1), loc="upper right")

    ax2.grid(which='minor', alpha=0.2)
    ax2.grid(which='major', alpha=0.4)
    ax1.grid(which='minor', alpha=0.2)
    ax1.grid(which='major', alpha=0.4)

def run_recruit_drop_rate_simulation():

    global fig
    global ax2
    global ax1

    fig, (ax1, ax2) = plt.subplots(2)
    fig.suptitle('Dynamic Recruit Drop Rates')
    fig.set_size_inches(12, 9)

    ani = FuncAnimation(
        fig,
        draw_recruit_drop_paths,
        frames=FRAMES,
        interval=100,
        repeat=False,
        init_func=init_plot,
    )

    plt.subplots_adjust(left=0.08, right=0.7, top=0.9, bottom=0.1, hspace=0.4)
    plt.show()

run_recruit_drop_rate_simulation()


from collections import defaultdict
fileIn = open("aoc5input.txt", "r")
D1 = defaultdict(int)
D2 = defaultdict(int)
for line in fileIn:
    initialPoint,endPoint = line.split("->")
    x1,y1 = initialPoint.split(",")
    x2,y2 = endPoint.split(",")
    x1 = int(x1)
    x2 = int(x2)
    y1 = int(y1)
    y2 = int(y2)
    changeX = x2 - x1
    changeY = y2 - y1
    for i in range(max(abs(changeX), abs(changeY))+1):
        if (changeX>0):
            x = x1 + i
        else:
            if (changeX<0):
                x = x1 - i
            else:
                x = x1
        if changeY > 0:
            y = y1 + i
        else:
            if changeY<0:
                y = y1 - i
            else:
                y = y1
        if changeX == 0 or changeY == 0:
            D1[(x,y)] += 1
        D2[(x,y)] += 1
numNonDiag = 0
numDiagInc = 0
fileIn.close()
for val in D1:
    if D1[val]>1:
        numNonDiag += 1
for val in D2:
    if D2[val]>1:
        numDiagInc += 1
print("part 1 and part 2, respectively:", numNonDiag, numDiagInc)


import matplotlib.pyplot as plt
import panel as pn

from plotnine import aes, element_rect, facet_wrap, geom_point, ggplot, stat_smooth, themes
from plotnine.data import mtcars
from dataviz_in_python import config

config.configure(url="lib_plotnine", title="Plotnine")

TEXT = """
# Plotnine: A Grammar of Graphics for Python

[Plotnine](https://plotnine.readthedocs.io/en/stable/) is an implementation of a grammar of graphics in Python, it is based on [ggplot2](https://ggplot2.tidyverse.org/). The grammar allows users to compose plots by explicitly mapping data to the visual objects that make up the plot. Its based on Matplotlib.

[Source Code](https://github.com/MarcSkovMadsen/dataviz-in-python/blob/main/src/dataviz_in_python/presentation/lib_plotnine.py)
"""
pn.panel(TEXT, css_classes=[config.TEXT_CLASS]).servable()

def get_plot(theme="default"):
    plt.style.use("default")
    if theme == "dark":
        plotnine_theme = themes.theme_dark() + themes.theme(
            plot_background=element_rect(fill="black", alpha=0)
        )
    else:
        plotnine_theme = themes.theme_xkcd()

    plot = (
        (
            ggplot(mtcars, aes("wt", "mpg", color="factor(gear)"))
            + geom_point()
            + stat_smooth(method="lm")
            + facet_wrap("~gear")
        )
        + plotnine_theme
        + themes.theme(figure_size=(16, 8))
    )
    return plot.draw()

plot = get_plot(theme=config.get_theme())

pn.pane.Matplotlib(plot, height=700, sizing_mode="stretch_both").servable()