#!/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


from flask import Flask, render_template
from log_analysis.__main__ import answer_all
from log_analysis.database import DBConnection
app = Flask(__name__)

@app.route('/')
@app.route('/answers')
def answers():
    with DBConnection() as conn:
        answers = answer_all(conn)
    return render_template('answers.html', answers=answers)

if __name__ == '__main__':
    app.debug = True
    app.run(host='0.0.0.0', port=5000)


import pickle

import requests
from bs4 import BeautifulSoup

from beer import Beer
from beerid import beerid
from storeid import storeid

if __name__ == '__main__':
    beerstore_url = 'https://www.thebeerstore.ca/beers/'
    stores = list()
    beers = list()
    out_msg: str

    # Search and select store
    with open('data/store_ids.pkl', 'rb') as infile:
        stores = pickle.load(infile)

    search_store_str = input('Search for preferred store: ')
    searched_stores = dict()

    counter = 1
    store: storeid
    for store in stores:
        if search_store_str.lower() in store.name.lower() or search_store_str.lower() in store.address.lower():
            searched_stores.update({counter: store})
            print(f'{counter}. {store.address}')
            counter += 1

    selected_store_index = int(input('Select a store by number: '))
    selected_store: storeid = searched_stores.get(selected_store_index)

    print(f'Store {selected_store.name}, at {selected_store.address} selected!\n')
    ##############################################################

    # Search and select beer
    with open('data/beer_ids.pkl', 'rb') as infile:
        beers = pickle.load(infile)

    searched_beer_str = input('Search for a beer: ')
    searched_beers = dict()

    counter = 1
    beerval: beerid
    for beerval in beers:
        if searched_beer_str.lower() in beerval.name.lower() or searched_beer_str.lower() in beerval.brewer.lower():
            searched_beers.update({counter: beerval})
            print(f'{counter}. {beerval.name}')
            counter += 1

    searched_beers_index = int(input('Select a beer by number: '))
    print(searched_beers_index)
    selected_beer: beerid = searched_beers.get(searched_beers_index)

    print(f'Beer {selected_beer.name} selected!\n')
    ##############################################################

    headers = {
        'User-Agent': '',
        'Cookie': f'beer_home_store={selected_store.id}',
    }

    request = requests.get(beerstore_url + selected_beer.link, headers=headers)
    status = request.status_code
    soup = BeautifulSoup(request.content, 'html.parser')

    beer_name = soup.find('div', class_='desc').find('h1', class_='capitalize').contents[0]

    html_beer_formats = soup.findAll('div', class_='total_cans')
    html_all_beer = []
    all_beer = []

    for beer_type in html_beer_formats:
        for beer in beer_type.findAll('li', class_='d-column d-row option _cart'):
            html_all_beer.append(beer)

    for beer in html_all_beer:
        str_quan_type = str(beer.find('div', class_='col_1 first col_same').contents[2]).strip().split(" ")
        quantity = int(str_quan_type[0])
        type = str_quan_type[2]
        serving = str_quan_type[3]

        stock = 0
        if (len(beer.find('div', class_='col_2 second col_same').contents) == 3):
            stock = int(str(beer.find('div', class_='col_2 second col_same').contents[2]).strip())
        elif (beer.find('div', class_='col_2 second col_same').find('span', class_='outStock')):
            stock = 0
        else:
            stock = 'Packup'

        sale_price = float(str(beer.find('span', class_='sale_price').contents[0]).strip()[1::])
        on_sale = True if beer.find('span', class_='sale_badge') else False

        all_beer.append(Beer(beer_name, quantity, type, serving, sale_price, stock, on_sale))

    all_beer.sort(reverse=True)

    # for beer in all_beer:
    #     print(f'\n{beer}')

    bestOption: Beer
    beer: Beer
    for beer in all_beer:
        if (beer.inStock()):
            bestOption = beer
            break

    try:
        bestOption
    except NameError:
        out_msg = f'There was no {beer_name} in stock, looks like you\'ll have to try another beer.'
    else:
        out_msg = (
            'Best Option:\n'
            f'{bestOption}'
        )

    print(out_msg)


from confluent_kafka import Producer

# Kafka broker configuration
bootstrap_servers = 'localhost:9094'
topic = 'test-topic'

# Test message to be sent
message = 'Hello, Kafka!'

def produce_message():
    # Create Kafka producer
    producer = Producer({'bootstrap.servers': bootstrap_servers})

    # Produce message to Kafka topic
    producer.produce(topic, message.encode())

    producer.flush()

    # Close the producer
    producer.close()

if __name__ == '__main__':
    # Produce a message
    produce_message()


import os
import copy
import logging

from typing import Optional

from lxml import etree  # type: ignore
from scielo_v3_manager.pid_manager import Manager
from scielo_v3_manager.v3_gen import generates

from prodtools.utils import fs_utils
from prodtools.utils import xml_utils

LOGGER = logging.getLogger(__name__)

class PidManagerExceedsIntentTimesError(Exception):
    ...

def update_xml_file(file_path, pids_to_append_in_xml):
    if not pids_to_append_in_xml:
        # nada para atualizar
        return None
    if not file_path:
        LOGGER.debug("Could not find XML path")
        return None
    try:
        tree = xml_utils.get_xml_object(file_path)
    except xml_utils.etree.XMLSyntaxError:
        LOGGER.info("%s is not a valid XML", file_path)
    else:
        # atualizar XML com os `pids_to_append_in_xml`
        _tree = add_article_id_to_etree(tree, pids_to_append_in_xml)
        write_etree_to_file(_tree, file_path)

def add_article_id_to_received_documents(
    pid_manager_info: dict,
    issn_id: str,
    year_and_order: str,
    received_docs: dict,
    documents_in_isis: dict,
    file_paths: dict,
    update_article_with_aop_status: callable,
) -> None:

    exceptions = []
    registered_v3_items = {}
    total = len(received_docs)
    times = 0
    MAX_TRIES = total
    results = {}
    while True:
        try:
            pid_manager = Manager(
                pid_manager_info['name'], pid_manager_info['timeout'])
            _add_article_id_to_received_documents(
                pid_manager,
                issn_id,
                year_and_order,
                received_docs,
                documents_in_isis,
                file_paths,
                update_article_with_aop_status,
                registered_v3_items,
                results,
            )
        except Exception as e:
            exceptions.append(str(e))
        finally:
            done = len([v for v in registered_v3_items.values() if v])
            if done == total:
                break
            times += 1
            if times > MAX_TRIES or len(exceptions) > 0:
                raise PidManagerExceedsIntentTimesError(
                    "Pid Manager failed to set v3 to %i documents. "
                    "Tried %i times. "
                    "Done for %i documents. %s %s" %
                    (total, times, done, results, "\n".join(exceptions))
                )

def _add_article_id_to_received_documents(
    pid_manager: Manager,
    issn_id: str,
    year_and_order: str,
    received_docs: dict,
    documents_in_isis: dict,
    file_paths: dict,
    update_article_with_aop_status: callable,
    registered_v3_items: dict,
    results: dict,
) -> None:
    """Atualiza article-id (scielo-v2 e scielo-v3) dos documentos recebidos.

    Params:
        pid_manager (PIDVersionsManager): instância de PIDVersionsManager para gerir pid da versão 3
        issn_id (str): ISSN do periódico
        year_and_order (str): Ano e ordem da issue processada
        received_docs (dict): Pacote de documentos recebidos para processar
        documents_in_isis (dict): Documentos já registrados na base isis (acron/volnum)
        file_paths (dict): arquivos do received_docs
        update_article_with_aop_status (callable): Função que recupera o AOP PID e modifica o
            artigo com este dado

    Returns:
        None
    """
    for xml_name, article in received_docs.items():
        if registered_v3_items.get(xml_name):
            continue
        file_path = file_paths.get(xml_name)
        if not file_path:
            LOGGER.debug("Could not find XML path for '%s' xml.", xml_name)

        pids_to_append_in_xml = []

        # Obtém v2 do XML
        pid_v2 = article.get_scielo_pid("v2")
        if pid_v2 is None:
            # se v2 não está presente no XML, gerar a partir dos metadados
            pid_v2 = get_scielo_pid_v2(issn_id, year_and_order, article.order)
            # anotar para ser inserido no XML
            pids_to_append_in_xml.append((pid_v2, "scielo-v2"))

        # Obtém v3 do XML
        pid_v3 = article.get_scielo_pid("v3")

        # Obtém previous do XML
        prev_pid = article.previous_article_pid
        if not prev_pid:
            # Não há previous do XML
            # Obtém previous_pid registrado na base ahead do artigo
            if update_article_with_aop_status:
                update_article_with_aop_status(article)
            # acessa previous_pid com `article.registered_aop_pid`
            prev_pid = article.registered_aop_pid
            if prev_pid:
                pids_to_append_in_xml.append((prev_pid, "previous-pid"))

        # consulta / registra / atualiza os dados na base pid_manager
        result = pid_manager.manage(
            v2=pid_v2, v3=pid_v3, aop=prev_pid,
            filename=os.path.basename(file_path),
            doi=article.doi,
            status="",
            generate_v3=generates)
        results[xml_name] = result

        record = result.get("saved") or result.get("registered") or {}
        v3 = record.get("v3")
        if v3:
            registered_v3_items[xml_name] = v3

        if pid_v3 is None:
            # se v3 não está no presente no XML
            article.registered_scielo_id = v3
            # anotar para ser inserido no XML
            pids_to_append_in_xml.append((v3, "scielo-v3"))

        # atualiza aop pid, se aplicável
        if record and not prev_pid:
            recovered_aop_pid = record.get("aop")
            if not recovered_aop_pid:
                if record.get("v2") and pid_v2 and pid_v2 != record.get("v2"):
                    recovered_aop_pid = record.get("v2")
            if recovered_aop_pid:
                pids_to_append_in_xml.append(
                    (recovered_aop_pid, "previous-pid"))

        # atualizar o XML com pids_to_append_in_xml
        update_xml_file(file_path, pids_to_append_in_xml)

def get_scielo_pid_v2(issn_id, year_and_order, order_in_issue):
    year = year_and_order[:4]
    order_in_year = year_and_order[4:].zfill(4)
    return "".join(("S", issn_id, year, order_in_year, order_in_issue))

def add_article_id_to_etree(
    tree: etree.ElementTree, pid_and_specific_use_items: list
) -> Optional[etree.ElementTree]:
    """Adiciona os pids v2 e v3 a árvore lxml de um documento"""
    if (
        tree is None
        or pid_and_specific_use_items is None
        or len(pid_and_specific_use_items) == 0
    ):
        return None

    _tree = copy.deepcopy(tree)

    article_meta = _tree.find(".//article-meta")

    if article_meta is None:
        LOGGER.debug(
            "Could not insert articles ids because the article-meta isn't found"
        )
        return None

    for id_value, specific_use in pid_and_specific_use_items:
        article_id = etree.Element("article-id")
        article_id.text = id_value
        article_id.set("specific-use", specific_use)
        article_id.set("pub-id-type", "publisher-id")
        article_meta.insert(0, article_id)

    return _tree

def write_etree_to_file(tree: etree.ElementTree, path: str) -> None:
    """Escreve uma árvore lxml em um arquivo de destino. Também
    garante que as entidades não serão modificadas por meio da função
    xml_utils.tostring(etree)."""

    if tree is None or path is None:
        return None

    fs_utils.write_file(path, xml_utils.tostring(tree))


from room import Room
from random import randint

class World:

	def __init__(self):

		self.currRoom = 0

		self.roomNames = ['living', 'garden', 'kitchen', 'bedroom']
		self.rooms = []
		desc = ["This room has a couch, chairs and TV.",
				"You have entered the living room. You can watch TV here",
				"This room has two sofas, chairs and a chandelier."]
		nei = [-1, 1, -1, 3]
		obj = "tv"
		act = "watch"
		self.rooms.append(Room(0, self.roomNames[0], desc, nei, obj, act))

		desc = ["This space has a swing, flowers and trees.",
				"You have arrived at the garden. You can exercise here",
				"This area has plants, grass and rabbits."]
		nei = [-1, -1, 0, 2]
		obj = "bike"
		act = "exercise"
		self.rooms.append(Room(1, self.roomNames[1], desc, nei, obj, act))

		desc = ["This room has a fridge, oven, and a sink.",
				"You have arrived in the kitchen. You can find food and drinks here.",
				"This living area has pizza, coke, and icecream."]
		nei = [1, -1, 3, -1]
		obj = "apple"
		act = "eat"
		self.rooms.append(Room(2, self.roomNames[2], desc, nei, obj, act))

		desc = ["This area has a bed, desk and a dresser.",
				"You have arrived in the bedroom. You can rest here.",
				"You see a wooden cot and a mattress on top of it."]
		nei = [0, 2, -1, -1]
		obj = "bed"
		act = "sleep"
		self.rooms.append(Room(3, self.roomNames[3], desc, nei, obj, act))

	def move(self, act, obj):
		if act == 'go':
			ind = 0
			if obj == 'north':
				ind = 0
			elif obj == 'east':
				ind = 1
			elif obj == 'west':
				ind = 2
			elif obj == 'south':
				ind = 3
			else:
				return 'not available'

			nextRoom = self.rooms[self.currRoom].neighbors[ind]
			if nextRoom == -1:
				return "not available"
			else:
				self.currRoom = nextRoom
			return self.desc()

		elif act == self.rooms[self.currRoom].action:
			if obj == self.rooms[self.currRoom].object:
				ans = ""
				if obj == "bed":
					ans = "You slept on bed."
				elif obj == "apple":
					ans = "You consumed apple."
				elif obj == "tv":
					ans = "You watched tv."
				else:
					ans = "You exercised with bike."
				return ans + '\nREWARD_' + self.rooms[self.currRoom].action + ' 1.'
			else:
				return 'not available'
		else:
			return 'not available'

	def desc(self):
		ind = randint(0,2)
		return self.rooms[self.currRoom].desc[ind]

# world = World()

# while True:
# 	act = raw_input("Action: ")
# 	obj = raw_input("Object: ")
# 	print(world.move(act, obj))
# 	print(world.desc())


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

import sys
import os
from PyQt4 import QtGui

class Test(QtGui.QWidget):

    def __init__(self):
        super(Test, self).__init__()

        self.initUI()

    def initUI(self):
        pixmap = QtGui.QPixmap("/Users/lukasfischer/github/in_vivo/MTH3/Figures/LF170110_2_Day201748_1_roi_6.png")

        print(pixmap.isNull())

        lbl = QtGui.QLabel(self)
        lbl.setPixmap(pixmap)

        self.move(300, 200)
        self.setWindowTitle('Test')
        self.show()

print(os.getcwd())
app = QtGui.QApplication(sys.argv)
ex = Test()
sys.exit(app.exec_())


# coding=utf-8
import re
import consts
import sundry as s
import sys
import linstordb
from execute.lvm import LVM
import subprocess

import time

class TimeoutError(Exception):
    pass

def judge_result(result):
    # 对命令进行结果根据正则匹配进行分类
    re_suc = re.compile('SUCCESS')
    re_war = re.compile('WARNING')
    re_err = re.compile('ERROR')
    """
    suc : 0
    suc with war : 1
    war : 2
    err : 3
    """
    if re_err.search(result):
        result = {'sts':3,'rst':result}
    elif re_suc.search(result) and re_war.search(result):
        messege_war = get_war_mes(result)
        result = {'sts':1,'rst':messege_war}
    elif re_suc.search(result):
        result = {'sts':0,'rst':result}
    elif re_war.search(result):
        messege_war = get_war_mes(result)
        result = {'sts':2,'rst':messege_war}

    if result:
        return result
    else:
        s.handle_exception()

def get_err_detailes(result):
    re_ = re.compile(r'Description:\n[\t\s]*(.*)\n')
    if re_.search(result):
        return (re_.search(result).group(1))

def get_war_mes(result):
    re_ = re.compile(r'\x1b\[1;33mWARNING:\n\x1b\[0m(.*)',re.DOTALL)
    if re_.search(result):
        return (re_.search(result).group(1))

@s.deco_cmd('linstor')
def execute_linstor_cmd(cmd,timeout=60):
    p = subprocess.Popen(cmd, stderr=subprocess.STDOUT, stdout=subprocess.PIPE, shell=True)
    t_beginning = time.time()
    seconds_passed = 0
    while True:
        if p.poll() is not None:
            break
        seconds_passed = time.time() - t_beginning
        if timeout and seconds_passed > timeout:
            p.terminate()
            raise TimeoutError(cmd, timeout)
        time.sleep(0.1)
    output = p.stdout.read().decode()
    return judge_result(output)

class Node():
    def __init__(self):
        pass

    """node操作"""
    # 创建集群节点
    def create_node(self, node, ip, nt):
        cmd = f'linstor node create {node} {ip}  --node-type {nt}'
        nt_value = [
            'Combined',
            'combined',
            'Controller',
            'Auxiliary',
            'Satellite']
        if nt not in nt_value:
            print('node type error,choose from ''Combined',
                  'Controller', 'Auxiliary', 'Satellite''')
        else:
            result = execute_linstor_cmd(cmd)
            if result['sts'] == 0:
                s.prt_log('SUCCESS', 0)
                return {'result':'SUCCESS','info':''}
            elif result['sts'] == 1:
                s.prt_log(f"SUCCESS\n{result['rst']}", 1)
                return {'result':'SUCCESS','info':result['rst']}
            elif result['sts'] == 2:
                s.prt_log(f"FAIL\n{result['rst']}", 1)
                return {'result':'FAIL','info':result['rst']}
            else:
                s.prt_log(f"FAIL\n{result['rst']}", 1)
                return {'result':'FAIL','info':result['rst']}

    # 删除node
    def delete_node(self, node):
        cmd = f'linstor node delete {node}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
            return {'result':'SUCCESS','info':''}
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
            return {'result':'SUCCESS','info':result['rst']}
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}

    def show_all_node(self, no_color='no'):
        collecter = linstordb.CollectData()
        if no_color == 'no':
            data = s.deco_color(collecter.get_all_node)()
        else:
            data = collecter.get_all_node()
        header = ["node", "node type", "res num", "stp num", "addr", "status"]
        table = s.make_table(header, data)
        s.prt_log(table,0)

    def show_one_node(self, node, no_color='no'):
        collecter = linstordb.CollectData()
        node_data = collecter.get_node_info(node)
        if not node_data:
            s.prt_log('The node does not exist',1)
        else:
            info = "node:%s\nnodetype:%s\nresource num:%s\nstoragepool num:%s\naddr:%s\nstatus:%s" % node_data
            if no_color == 'no':
                data_node = s.deco_color(collecter.get_one_node)(node)
                data_stp = s.deco_color(collecter.get_sp_in_node)(node)
            else:
                data_node = collecter.get_one_node(node)
                data_stp = collecter.get_sp_in_node(node)

            header_node = ['res_name', 'stp_name', 'size', 'device_name', 'used', 'status']
            header_stp = ['stp_name', 'node_name', 'res_num', 'driver', 'pool_name', 'free_size', 'total_size', 'snapshots',
                          'status']
            table_node = s.make_table(header_node, data_node)
            table_stp = s.make_table(header_stp, data_stp)
            result = '\n'.join([info, str(table_node), str(table_stp)])
            s.prt_log(result,0)

class StoragePool():
    def __init__(self):
        pass

    """storagepool操作"""
    # 创建storagepool
    def create_storagepool_lvm(self, node, stp, vg):
        obj_lvm = LVM()
        if not obj_lvm.is_vg_exists(vg):
            s.prt_log(f'Volume group:"{vg}" does not exist',1)
            return
        cmd = f'linstor storage-pool create lvm {node} {stp} {vg}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
            return {'result':'SUCCESS','info':''}
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
            return {'result':'SUCCESS','info':result['rst']}
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}

    def create_storagepool_thinlv(self, node, stp, tlv):
        obj_lvm = LVM()
        if not obj_lvm.is_thinlv_exists(tlv):
            s.prt_log(f'Thin logical volume:"{tlv}" does not exist',1)
            return
        cmd = f'linstor storage-pool create lvmthin {node} {stp} {tlv}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
            return {'result':'SUCCESS','info':''}
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
            return {'result':'SUCCESS','info':result['rst']}
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}

    # 删除storagepool -- ok
    def delete_storagepool(self, node, stp):
        cmd = f'linstor storage-pool delete {node} {stp}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
            return {'result':'SUCCESS','info':''}
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
            return {'result':'SUCCESS','info':result['rst']}
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return {'result':'FAIL','info':result['rst']}

    def show_all_sp(self,no_color='no'):
        collector = linstordb.CollectData()
        if no_color == 'no':
            data = s.deco_color(collector.get_all_sp)()
        else:
            data = collector.get_all_sp()
        header = ['stp_name','node_name','res_num','driver','pool_name','free_size','total_size','snapshots','status']
        table = s.make_table(header, data)
        s.prt_log(table,0)

    def show_one_sp(self,sp,no_color='no'):
        collector = linstordb.CollectData()
        info_data = collector.get_sp_info(sp)
        if info_data[0] == 0:
            s.prt_log('The storagepool does not exist',1)
        else:
            info = 'The storagepool name for %s nodes is %s,they are %s.'%(info_data)
            if no_color == 'no':
                data = s.deco_color(collector.get_one_sp)(sp)
            else:
                data = collector.get_one_sp(sp)
            header = ['res_name', 'size', 'device_name', 'used', 'status']
            table = s.make_table(header, data)
            result = '\n'.join([info, str(table)])
            s.prt_log(result,0)

class Resource():
    def __init__(self):
        pass

    def collect_args(self,node,sp):
        """
        收集输入的参数，进行处理
        :param node: 列表，node名
        :param sp: 列表，storagepool名
        :return: 字典
        """

        dict_args = {}
        if len(sp) == 1:
            for node_one in node:
                dict_args.update({node_one:sp[0]})
        else:
            for node_one,sp_one in zip(node,sp):
                dict_args.update({node_one:sp_one})
        return dict_args

    def execute_create_res(self,res,node,sp):
        # 执行在指定节点和存储池上创建resource
        # 成功返回空字典，失败返回 {节点：错误原因}
        cmd = f'linstor resource create {node} {res} --storage-pool {sp}'
        try:
            result = execute_linstor_cmd(cmd)
            if result['sts'] == 2:
                s.prt_log(get_war_mes(result),1)
            if result['sts'] == 0:
                s.prt_log(f'Resource {res} was successfully created on Node {node}',0)
                return {}
            elif result['sts'] == 1:
                s.prt_log(f'Resource {res} was successfully created on Node {node}',0)
                return {}
            elif result['sts'] == 3:
                fail_cause = get_err_detailes(result['rst'])
                dict_fail = {node: fail_cause}
                return dict_fail

        except TimeoutError:
            result = f'{res} created timeout on node {node}, the operation has been cancelled'
            s.prt_log(result,2)
            return {node:'Execution creation timeout'}

    # 创建resource相关
    def linstor_delete_rd(self, res):
        cmd = f'linstor rd d {res}'
        execute_linstor_cmd(cmd)

    def linstor_create_rd(self, res):
        cmd = f'linstor rd c {res}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            return True
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return result
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            return result

    def linstor_create_vd(self, res, size):
        cmd = f'linstor vd c {res} {size}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            return True
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            self.linstor_delete_rd(res)
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
            self.linstor_delete_rd(res)
            return result

    # 创建resource 自动
    def create_res_auto(self, res, size, num):
        cmd = f'linstor r c {res} --auto-place {num}'
        if self.linstor_create_rd(res) is True and self.linstor_create_vd(res, size) is True:
            result = execute_linstor_cmd(cmd)
            if result['sts'] == 0:
                s.prt_log('SUCCESS', 0)
                return True
            elif result['sts'] == 1:
                s.prt_log(f"SUCCESS\n{result['rst']}", 1)
                return True
            elif result['sts'] == 2:
                self.linstor_delete_rd(res)
                s.prt_log(f"FAIL\n{result['rst']}", 1)
                return result
            else:
                self.linstor_delete_rd(res)
                s.prt_log(f"FAIL\n{result['rst']}", 1)
                return result

        else:
            s.prt_log('The resource already exists',0)
            return ('The resource already exists')

    def create_res_manual(self, res, size, node, sp):
        dict_all_fail = {}
        dict_args = self.collect_args(node,sp)

        if self.linstor_create_rd(res) is True and self.linstor_create_vd(res, size) is True:
            pass
        else:
            s.prt_log('The resource already exists',1)
            return ('The resource already exists')

        for node_one,sp_one in dict_args.items():
            dict_one_result = self.execute_create_res(res,node_one,sp_one)
            dict_all_fail.update(dict_one_result)

        if len(dict_all_fail.keys()) == len(node):
            self.linstor_delete_rd(res)
        if len(dict_all_fail.keys()):
            fail_info = (f"Creation failure on {' '.join(dict_all_fail.keys())}\n")
            for node, cause in dict_all_fail.items():
                fail_cause = f"{node}:{cause}\n"
                fail_info = fail_info + fail_cause
            s.prt_log(fail_info,2)
            return dict_all_fail
        else:
            return True

    # 添加mirror（自动）
    def add_mirror_auto(self, res, num):
        cmd = f'linstor r c {res} --auto-place {num}'
        result = execute_linstor_cmd(cmd)
        # result = self.judge_result(output)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 2)

    def add_mirror_manual(self, res, node, sp):
        dict_all_fail = {}
        dict_args = self.collect_args(node,sp)

        for node_one,sp_one in dict_args.items():
            dict_one_result = self.execute_create_res(res,node_one,sp_one)
            dict_all_fail.update(dict_one_result)

        if len(dict_all_fail.keys()):
            fail_info = (f"Creation failure on {' '.join(dict_all_fail.keys())}\n")
            for node, cause in dict_all_fail.items():
                fail_cause = f"{node}:{cause}\n"
                fail_info = fail_info + fail_cause
            s.prt_log(fail_info,1)
            return dict_all_fail
        else:
            return True

    # 创建resource --diskless
    def create_res_diskless(self, node, res):
        cmd = f'linstor r c {node[0]} {res} --diskless'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 2)

    # 删除resource,指定节点
    def delete_resource_des(self, node, res):
        cmd = f'linstor resource delete {node} {res}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 2)

    # 删除resource，全部节点
    def delete_resource_all(self, res):
        cmd = f'linstor resource-definition delete {res}'
        result = execute_linstor_cmd(cmd)
        if result['sts'] == 0:
            s.prt_log('SUCCESS', 0)
        elif result['sts'] == 1:
            s.prt_log(f"SUCCESS\n{result['rst']}", 1)
        elif result['sts'] == 2:
            s.prt_log(f"FAIL\n{result['rst']}", 1)
        else:
            s.prt_log(f"FAIL\n{result['rst']}", 2)

    def show_all_res(self,no_color='no'):
        collecter = linstordb.CollectData()
        if no_color == 'no':
            data = s.deco_color(collecter.get_all_res)()
        else:
            data = collecter.get_all_res()
        header = ["resource", "mirror_way", "size", "device_name", "used"]
        table = s.make_table(header,data)
        s.prt_log(table,0)

    def show_one_res(self,res,no_color='no'):
        collecter = linstordb.CollectData()
        res_data = collecter.get_res_info(res)
        if not res_data:
            s.prt_log('The resource does not exist',1)
        else:
            info = ("resource:%s\nmirror_way:%s\nsize:%s\ndevice_name:%s\nused:%s" %res_data)
            if no_color == 'no':
                data = s.deco_color(collecter.get_one_res)(res)
            else:
                data = collecter.get_one_res(res)
            header = ['node_name', 'stp_name', 'drbd_role', 'status']
            table = s.make_table(header,data)
            result = '\n'.join([info, str(table)])
            s.prt_log(result,0)


#!/usr/bin/env python3

from numpy import empty
import rosbag
import rospy
import datetime

from file_info import FileInfo

class RosBagHandler:
  position_topics = {'/gps':'DriX',
    '/project11/mesobot/nav/position':'Mesobot',
    '/project11/nui/nav/position':'nui',
    '/project11/nautilus/position':'Nautilus',
    '/mothership_gps':'Mothership'
  }

  def __init__(self):
    pass

  def needsProcessing(self, file: FileInfo):
    if file.local_path.suffix == '.bag':
      if file.has_meta_value(self, 'start_time') and not file.is_modified():
        return False
      if file.meta is not None and 'RosBagIndexHandler' in file.meta:
        if 'indexed' in file.meta['RosBagIndexHandler'] and not file.meta['RosBagIndexHandler']['indexed']:
          return False
      if file.has_meta_value(self, 'message_count') and file.get_meta_value(self, 'message_count') == 0:
        return False
      return True
    return False

  def process(self, file: FileInfo) -> FileInfo:
    try:
      bag = rosbag.Bag(file.source_path())
    except Exception as e:
      print("error opening bag file",file.local_path,e)
      print(type(e))
      return
    try:
      file.update_meta_value(self, 'message_count', bag.get_message_count())
      if bag.get_message_count() == 0:
        return
      file.update_meta_value(self, 'start_time', bag.get_start_time())
      file.update_meta_value(self, 'end_time', bag.get_end_time())
      tt = bag.get_type_and_topic_info()
      topics = []
      for t in RosBagHandler.position_topics:
        if t in tt.topics:
          topics.append(t)
      if len(topics) == 0:
        return
    except Exception as e:
      print("error getting times from bag file",file.local_path,e)
      print(type(e))
      return

    tracks = {}
    last_report_times = {}
    interval = rospy.Duration(secs=1.0)
    try:
      for topic, msg, t in bag.read_messages(topics=topics):
        vehicle = RosBagHandler.position_topics[topic]
        if not vehicle in tracks:
          tracks[vehicle] = []
          last_report_times[vehicle] = None
        if topic in ('/gps','/mothership_gps'):
          if msg.fix_quality > 0:
            if last_report_times[vehicle] is None or msg.header.stamp - last_report_times[vehicle] >= interval:
              fix = {'timestamp': msg.header.stamp.to_sec()}
              fix['latitude'] = msg.latitude
              fix['longitude'] = msg.longitude
              fix['altitude'] = 0.0
              tracks[vehicle].append(fix)
              last_report_times[vehicle] = msg.header.stamp
        elif topic in ('/project11/mesobot/nav/position','/project11/nui/nav/position'):
          if last_report_times[vehicle] is None or msg.header.stamp - last_report_times[vehicle] >= interval:
              fix = {'timestamp': msg.header.stamp.to_sec()}
              fix['latitude'] = msg.pose.position.latitude
              fix['longitude'] = msg.pose.position.longitude
              fix['altitude'] = msg.pose.position.altitude
              tracks[vehicle].append(fix)
              last_report_times[vehicle] = msg.header.stamp
        # elif topic in ('/project11/nui/nav/position',):
        #   if last_report_times[vehicle] is None or msg.header.stamp - last_report_times[vehicle] >= interval:
        #       fix = {'timestamp': msg.header.stamp.to_sec()}
        #       fix['latitude'] = msg.position.latitude
        #       fix['longitude'] = msg.position.longitude
        #       fix['altitude'] = msg.position.altitude
        #       tracks[vehicle].append(fix)
        #       last_report_times[vehicle] = msg.header.stamp
        else:
          if msg.status.status >= 0:
            if last_report_times[vehicle] is None or msg.header.stamp - last_report_times[vehicle] >= interval:
              fix = {'timestamp': msg.header.stamp.to_sec()}
              fix['latitude'] = msg.latitude
              fix['longitude'] = msg.longitude
              fix['altitude'] = msg.altitude
              tracks[vehicle].append(fix)
              last_report_times[vehicle] = msg.header.stamp

    except Exception as e:
      print("error extracting nav from bag file",file.local_path, e)

    bounds = {}
    tracks_for_meta = {}
    for v in tracks:
      if len(tracks[v]):
        min_lat = max_lat = tracks[v][0]['latitude']
        min_lon = max_lon = tracks[v][0]['longitude']
        for fix in tracks[v]:
          min_lat = min(min_lat, fix['latitude'])
          max_lat = max(max_lat, fix['latitude'])
          min_lon = min(min_lon, fix['longitude'])
          max_lon = max(max_lon, fix['longitude'])
        bounds[v] = {'min': {'latitude': min_lat, 'longitude': min_lon}, 'max': {'latitude': max_lat, 'longitude': max_lon}}
        tracks_for_meta[v] = tracks[v]

    if len(bounds):
      file.update_meta_value(self, 'bounds', bounds)
      file.update_meta_value(self, 'tracks', tracks_for_meta)


from qtpy import QtWidgets

from hyperspyui.widgets.signallist import SignalList

class PickXSignalsWidget(QtWidgets.QWidget):

    def __init__(self, signal_wrappers, x, parent=None, titles=None,
                 wrap_col=4):
        super().__init__(parent)
        grid = QtWidgets.QGridLayout()
        self.pickers = []
        self._x = x
        self._sws = signal_wrappers

        if titles is None:
            titles = list(("",) * x)
        elif isinstance(titles, str):
            titles = list((titles,) * x)
        elif len(titles) < x:
            titles.extend(list(("",) * (x-len(titles))))

        for i in range(x):
            picker = SignalList(signal_wrappers, self, False)
            grid.addWidget(QtWidgets.QLabel(titles[i]), 2*(i // wrap_col),
                           i % wrap_col)
            grid.addWidget(picker, 1 + 2*(i // wrap_col), i % wrap_col)
            self.pickers.append(picker)
        self.setLayout(grid)

    def get_selected(self):
        signals = []
        for i in range(self._x):
            signals.append(self.pickers[i].get_selected())
        if self._x == 1:
            return signals[0]
        return signals

    def unbind(self):
        for i in range(self._x):
            self.pickers[i].unbind(self._sws)


from flask import Flask, jsonify, request
from flask_socketio import SocketIO, emit
from pymongo import MongoClient
from bson.objectid import ObjectId
import subprocess
import sys
import os

# def run_cli():
#     cli_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'cli.py')
#     subprocess.Popen([sys.executable, cli_script])

app = Flask(__name__)
socketio = SocketIO(app)

# Connect to MongoDB
client = MongoClient('mongodb://localhost:27017/')
db = client['my_database']
collection = db['users']

@app.route('/')
def index():
    return "Welcome to the Firebase!"

@app.route('/users', methods=['GET', 'POST'])
def users():
    if request.method == 'GET':
        cursor = collection.find()
        users = [user for user in cursor]
        return jsonify(users)
    elif request.method == 'POST':
        user_data = request.json
        result = collection.insert_one(user_data)
        socketio.emit('data_updated', {'event': 'created', 'data': user_data})
        return jsonify({"success": True, "message": "User created", "_id": str(result.inserted_id)})

@app.route('/users/<int:user_id>', methods=['PUT', 'PATCH', 'DELETE'])
def user(user_id):
    print(f"Request received: method={request.method}, user_id={user_id}")
    #oid = ObjectId(user_id)
    oid = user_id
    if request.method == 'PUT':
        user_data = request.json
        result = collection.replace_one({'_id': oid}, user_data)
        if result.matched_count == 0:
            socketio.emit('data_updated', {'event': 'created', 'data': user_data},)
            message = "User created"
        else:
            socketio.emit('data_updated', {'event': 'updated', 'data': user_data})
            message = "User updated"
        return jsonify({"success": True, "message": message})
    elif request.method == 'PATCH':
        update_data = request.json
        result = collection.update_one({'_id': oid}, {'$set': update_data})
        print(f"Query filter: {{'_id': {user_id}}}")  # Debug print
        print(f"Result: {result.raw_result}")  # Debug print

        if result.modified_count == 1:
            #socketio.emit('data_updated', {'event': 'updated', 'data': update_data}, broadcast=True)
            socketio.emit('data_updated', {'event': 'updated', 'data': update_data})
            return jsonify({"success": True, "message": "User updated"})
        else:
            return jsonify({"success": False, "message": "User not found"})
    elif request.method == 'DELETE':
        result = collection.delete_one({'_id': oid})
        if result.deleted_count == 1:
            #socketio.emit('data_updated', {'event': 'deleted', 'user_id': user_id}, broadcast=True)
            socketio.emit('data_updated', {'event': 'deleted', 'user_id': user_id})
            return jsonify({"success": True, "message": "User deleted"})
        else:
            return jsonify({"success": False, "message": "User not found"})

if __name__ == '__main__':
    # run_cli()
    socketio.run(app, debug=True)


# -*- 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)))


#!/usr/bin/env python3
#
# Carson Carpenter
# CPSC 223P-03
# 2020-09-010
# carson.carpenter7@csu.fullerton.edu
#
"""
This is my program for converting and printing the string
representing the ASCII values of its binary representation
Character is inputed through the terminal
"""
import sys

def convert_from_binary():
    """ Converts binary to ASCII"""
    if len(sys.argv) > 1:
        character = sys.argv[1]
        for character in sys.argv[1:]:
            print(chr(int(character, 2)), end='')
        print()

if __name__ == '__main__':
    convert_from_binary()