#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 17 16:17:10 2020

@author: danbiderman
"""

"""In: time indices; DLC tracked pts in 2D (cam1+cam2); DLC 3D tracked pts; recovered dict
we loop over time points, and present the images along with our 3D plot
figure should be presented using """
import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objs as go
import cv2
from scipy.spatial.transform import Rotation as R
from matplotlib import colors
from PIL import Image

# this function assumes that we're plotting the first n frames of our dictionaries
def video_recover_plots(
    noisy_pose_dict,
    recovered_pose_dict,
    stacked_images,
    DLC_sliced1,
    DLC_sliced2,
    reproj_list_of_dicts,
    save_name,
    title,
    n_frames,
    lim_abs,
):
    """seems to work.this version has the gridspec including
    what i learned - using p3.Axes3D along with regular ax.plot was helpful.
    in ax.plot I just use x,y,z, coordinates.
    there is no init function like in the 2d plots, we initialize the lines
    with the vales at the zeroth timepoint"""
    from matplotlib.animation import FuncAnimation
    from matplotlib.gridspec import GridSpec

    fig = plt.figure(constrained_layout=True, figsize=(11, 6))

    gs = GridSpec(2, 2, figure=fig)
    ax1 = fig.add_subplot(gs[:2, 0])
    import mpl_toolkits.mplot3d.axes3d as p3

    # identical to ax1 = plt.subplot(gs.new_subplotspec((0, 0), colspan=3))
    ax2 = fig.add_subplot(
        gs[:2, 1], projection="3d"
    )  # https://matplotlib.org/3.1.1/gallery/mplot3d/subplot3d.html

    # fig = plt.figure()
    ax2.view_init(elev=-82.0, azim=-72)  # specific for costa data, found manually

    #    ax.set_xlim3d([-2.0, 2.0])
    #    ax.set_xlabel('X')
    #
    #    ax.set_ylim3d([-2.0, 2.0])
    #    ax.set_ylabel('Y')
    #
    #    ax.set_zlim3d([-2.0, 2.0])
    #    ax.set_zlabel('Z')

    im_stacked = ax1.imshow(
        stacked_images[:, :, 0], cmap="gray"
    )  # all the lines below consider a SPECIFIC CROP
    (line_1,) = ax1.plot(
        DLC_sliced1[0 + 1, [0, 3, 6]] - 200,
        DLC_sliced1[0 + 1, [1, 4, 7]],
        "r*",
        markersize=10,
    )

    (line_1_reproj,) = ax1.plot(
        reproj_list_of_dicts[0]["x_coords"][:, 0 + 1] - 200,
        reproj_list_of_dicts[0]["y_coords"][:, 0 + 1],
        "b+",
        markersize=10,
    )
    #    scatt_1 = ax1.scatter(DLC_sliced1[0+1,[0,3,6]]-200, DLC_sliced1[0+1,[1,4,7]],
    #        color='red', marker='*', s=50)
    (line_2,) = ax1.plot(
        DLC_sliced2[0 + 1, [0, 3, 6]] - 220 + 300,
        DLC_sliced2[0 + 1, [1, 4, 7]] - 50,  # +300 due to stack
        "r*",
        markersize=10,
    )
    (line_2_reproj,) = ax1.plot(
        reproj_list_of_dicts[1]["x_coords"][:, 0 + 1] - 220 + 300,
        reproj_list_of_dicts[1]["y_coords"][:, 0 + 1] - 50,
        "b+",
        markersize=10,
    )

    line_v = ax1.axvline(x=300, ymin=0, ymax=300, color="white", linewidth=5)
    ax1.set_title("View 1             View 2", fontsize=16, weight="bold")
    ax1.axis("off")
    ax1.legend(loc="lower left", labels=["DLC 2D", "BA-reproj"])

    # init lines with true data, no init method.
    (line_noise,) = ax2.plot(
        noisy_pose_dict["x_coords"][:, 0],
        noisy_pose_dict["y_coords"][:, 0],
        noisy_pose_dict["z_coords"][:, 0],
        "ro-",
        lw=4,
        markersize=5,
        label="DLC 3D",
    )
    (line_recovered,) = ax2.plot(
        recovered_pose_dict["x_coords"][:, 0],
        recovered_pose_dict["y_coords"][:, 0],
        recovered_pose_dict["z_coords"][:, 0],
        "ko-",
        lw=4,
        markersize=5,
        label="BA_3D",
        alpha=0.5,
    )
    #    plt.legend(loc = "lower left")
    #    plt.title(title, fontsize = 16)

    # dan 2/28 can be avoided
    # new for costa data
    #     ax2.set_xlim3d([-2.0, 2.0])
    #     #ax.set_xticks([])
    #     ax2.set_xticks([-2.0, -1.0, 0.0, 1.0, 2.0])
    #     ax2.set_xticklabels([])
    #     ax2.tick_params('x')
    #     ax2.set_xlabel('X', fontsize=12, labelpad = 0.1)
    #     ax2.set_ylim3d([-2.0, 2.0])
    #     ax2.set_yticks([-2.0, -1.0, 0.0, 1.0, 2.0])
    #     ax2.set_yticklabels([])
    #     ax2.set_ylabel('Y', fontsize=12)
    #     ax2.set_zlim3d([-2.0, 2.0])
    #     ax2.set_zticks([-2.0, -1.0, 0.0, 1.0, 2.0])
    #     ax2.set_zticklabels([])
    #     ax2.set_zlabel('Z', fontsize=12)
    ax2.legend(loc="lower left", fontsize=10)
    ax2.set_title("DLC Calib VS BA", fontsize=16, weight="bold", pad=18)

    # need to get images here
    def animate(
        i,
        stacked_images,
        im_stacked,
        DLC_sliced1,
        line_1,
        DLC_sliced2,
        line_2,
        reproj_list,
        line_1_reproj,
        line_2_reproj,
        line_v,
        noisy_pose_dict,
        line_noise,
        recovered_pose_dict,
        line_recovered,
    ):  # note, the loop has to see all these datasets

        im_stacked.set_data(
            stacked_images[:, :, i]
        )  # all the lines below consider a SPECIFIC CROP
        line_1.set_data(
            DLC_sliced1[i + 1, [0, 3, 6]] - 200, DLC_sliced1[i + 1, [1, 4, 7]]
        )
        line_1_reproj.set_data(
            reproj_list_of_dicts[0]["x_coords"][:, i + 1] - 200,
            reproj_list_of_dicts[0]["y_coords"][:, i + 1],
        )
        line_2.set_data(
            DLC_sliced2[i + 1, [0, 3, 6]] - 220 + 300,
            DLC_sliced2[i + 1, [1, 4, 7]] - 50,
        )
        line_2_reproj.set_data(
            reproj_list_of_dicts[1]["x_coords"][:, i + 1] - 220 + 300,
            reproj_list_of_dicts[1]["y_coords"][:, i + 1] - 50,
        )
        line_v = ax1.axvline(
            x=300, ymin=0, ymax=300, color="white", linewidth=5
        )  # .set_data(x=300, ymin = 0, ymax = 300)
        # noisy obs
        line_noise.set_data(
            noisy_pose_dict["x_coords"][:, i], noisy_pose_dict["y_coords"][:, i]
        )
        line_noise.set_3d_properties(noisy_pose_dict["z_coords"][:, i])

        # recovered
        line_recovered.set_data(
            recovered_pose_dict["x_coords"][:, i], recovered_pose_dict["y_coords"][:, i]
        )
        line_recovered.set_3d_properties(recovered_pose_dict["z_coords"][:, i])

        return (
            im_stacked,
            line_1,
            line_2,
            line_1_reproj,
            line_2_reproj,
            line_v,
            line_noise,
            line_recovered,
        )  # not sure the return is needed. check.

    anim = FuncAnimation(
        fig,
        animate,
        fargs=(
            stacked_images,
            im_stacked,
            DLC_sliced1,
            line_1,
            DLC_sliced2,
            line_2,
            reproj_list_of_dicts,
            line_1_reproj,
            line_2_reproj,
            line_v,
            noisy_pose_dict,
            line_noise,
            recovered_pose_dict,
            line_recovered,
        ),  # all inputs to loop
        frames=n_frames,
        interval=50,
        blit=False,
    )  # blit false important?

    # anim.save("recover.mp4")
    # anim.save(save_name + '.gif', writer='Imagekick') # works
    anim.save(save_name + ".gif", writer="pillow")  # pillow important

def plot_reproj_traces(points_2d, points_proj, frame_range, title, savename):
    """choose frame range according to num_frames_analyzed and
    num_body_parts and num_cameras. above, we take each body part coord, and stack
    it for all body parts and cameras."""
    plt.suptitle(title, fontsize=14)
    plt.subplot(211)
    plt.ylabel("x coord")
    plt.plot(points_2d[frame_range, 0], color="gray", linewidth=3)
    plt.plot(points_proj[frame_range, 0], color="red", linestyle="dashed", linewidth=1)
    plt.subplot(212)
    plt.ylabel("y coord")
    plt.plot(points_2d[frame_range, 1], color="gray", linewidth=3, label="DLC 2D")
    plt.plot(
        points_proj[frame_range, 1],
        color="red",
        linestyle="dashed",
        linewidth=1,
        label="re-projection",
    )
    plt.xlabel("frame num.")
    plt.legend(loc="lower right")
    plt.savefig(savename + ".png")

def plot_image_labels(
    image,
    coord_list_of_dicts,
    index,
    color_list,
    label_list=None,
    ax=None,
    top_img_height=0,
    pad=10,
):
    """ToDo: add multiple data sources, so instead of x_arr, y_arr have a list of dicts
    with x,y coords."""
    if len(image.shape) == 2:
        nrows, ncols = image.shape
    else:
        nrows, ncols = image.shape[0:2]

    assert len(color_list) == len(coord_list_of_dicts)

    ax.axis("off")
    ax.set_xlim(-pad, ncols + pad)  # pad x,y axes
    ax.set_ylim(nrows + pad, -pad)

    ax.imshow(image, "gray")
    for i in range(len(coord_list_of_dicts)):
        x_coord = np.copy(coord_list_of_dicts[i]["x_coords"][index, :])
        y_coord = np.copy(coord_list_of_dicts[i]["y_coords"][index, :])

        # indices 1 and 3 correspond to the bottom image.
        # Recall, we subtracted out the height of the top
        # image earlier.
        if i % 2 != 0:
            y_coord += top_img_height

        if label_list is not None:
            ax.scatter(
                x_coord, y_coord, color=color_list[i], label=label_list[i],
            )
        else:
            ax.scatter(x_coord, y_coord, color=color_list[i])

def plot_3d_points(coord_list_of_dicts, lims_dict, index, color_list, ax):
    """plot a single frame
    ToDo: specify plotting patams, allow for multiple inputs"""

    assert len(color_list) == len(coord_list_of_dicts)

    ax.set_xlim3d(lims_dict["x"])
    ax.set_ylim3d(lims_dict["y"])
    ax.set_zlim3d(lims_dict["z"])
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")

    for i in range(len(coord_list_of_dicts)):
        ax.plot(
            coord_list_of_dicts[i]["x_coords"][index, :],
            coord_list_of_dicts[i]["y_coords"][index, :],
            coord_list_of_dicts[i]["z_coords"][index, :],
            linestyle="None",
            marker="o",
            color=color_list[i],
            markersize=4,
            label=None,
        )

def vector_plot(
    tvects,
    is_vect=True,
    orig=[0, 0, 0],
    # names=["x", "y", "z"],
    names=["", "", ""],
    cam_name="",
    colors=["red", "blue", "green"],
):
    """Plot vectors using plotly"""

    if is_vect:
        if not hasattr(orig[0], "__iter__"):
            coords = [[orig, np.sum([orig, v], axis=0)] for v in tvects]
        else:
            coords = [[o, np.sum([o, v], axis=0)] for o, v in zip(orig, tvects)]
    else:
        coords = tvects

    data = []
    for i, c in enumerate(coords):
        X1, Y1, Z1 = zip(c[0])
        X2, Y2, Z2 = zip(c[1])
        vector = go.Scatter3d(
            x=[X1[0], X2[0]],
            y=[Y1[0], Y2[0]],
            z=[Z1[0], Z2[0]],
            text=[cam_name, names[i]],
            marker=dict(size=[0, 5], color="blue"),
            line=dict(width=5, color=colors[i]),
            mode="lines+markers+text",
            name=names[i] + cam_name,
        )
        data.append(vector)

    return data
    # fig = go.Figure(data=data,layout=layout)
    # fig.show()

def draw_circles(
    img, points, point_colors=None, point_size=10, marker_type=None, thickness=-1
):

    if point_colors is not None:
        if not isinstance(point_colors, list):
            point_colors = [point_colors] * points.shape[0]
        assert len(point_colors) == points.shape[0]

    for i, point in enumerate(points):
        if point_colors is not None:
            color = [i * 255 for i in colors.to_rgb(point_colors[i])]
        else:
            color = (255, 255, 255)

        if marker_type is not None:
            assert marker_type in ("cross", "diamond", "square")
            if marker_type == "cross":
                marker = cv2.MARKER_CROSS

            cv2.drawMarker(
                img,
                (point[0], point[1]),
                markerType=marker,
                markerSize=point_size,
                thickness=2,
                color=(color[2], color[1], color[0], 0),
            )
        else:
            cv2.circle(
                img,
                (point[0], point[1]),
                point_size,
                (color[2], color[1], color[0]),
                thickness,
            )

    return img

def slope(x1, y1, x2, y2):
    ###finding slope
    if x2 != x1:
        return (y2 - y1) / (x2 - x1)
    else:
        return "NA"

def drawLine(image, x1, y1, x2, y2, color=None):

    m = slope(x1, y1, x2, y2)
    h, w = image.shape[:2]
    if m != "NA":
        ### here we are essentially extending the line to x=0 and x=width
        ### and calculating the y associated with it
        ##starting point
        px = 0
        py = -(x1 - 0) * m + y1
        ##ending point
        qx = w
        qy = -(x2 - w) * m + y2
    else:
        ### if slope is zero, draw a line with x=x1 and y=0 and y=height
        px, py = x1, 0
        qx, qy = x1, h

    if color is not None:
        color = [i * 255 for i in colors.to_rgb(color)]
        image = cv2.line(
            image,
            (int(px), int(py)),
            (int(qx), int(qy)),
            (color[2], color[1], color[0]),
            2,
        )
    else:
        image = cv2.line(
            image, (int(px), int(py)), (int(qx), int(qy)), (255, 255, 255), 2
        )
    return image

def skew(vector):
    """
    this function returns a numpy array with the skew symmetric cross product matrix for vector.
    the skew symmetric cross product matrix is defined such that
    np.cross(a, b) = np.dot(skew(a), b)

    :param vector: An array like vector to create the skew symmetric cross product matrix for
    :return: A numpy array of the skew symmetric cross product vector
    """

    return np.array(
        [
            [0, -vector[2], vector[1]],
            [vector[2], 0, -vector[0]],
            [-vector[1], vector[0], 0],
        ]
    )

def plot_cams_and_points(
    cam_group=None,
    points_3d=None,
    title=None,
    scene_lims=None,
    point_size=2,
    scene_aspect=None,
    scene_camera=None,
    show_plot=True,
    skeleton_bp=None,
    skeleton_lines=None,
    line_colors=None,
    point_colors=None,
    legend=False,
    font_size=15,
    plot_names=False,
    cam_names=None
):
    """
    Plots the coordinate systems of the cameras along with the 3D points
    """
    if points_3d is not None:
        scatter_data = [
            go.Scatter3d(
                x=points_3d[:, 0],
                y=points_3d[:, 1],
                z=points_3d[:, 2],
                mode="markers",
                marker=dict(size=point_size, color=point_colors),
            )
        ]

        data = scatter_data
    else:
        data = []

    if cam_group is not None:
        for i, cam in enumerate(cam_group.cameras):
            rot_vec = R.from_rotvec(cam.get_rotation())
            R_mat = rot_vec.as_matrix() * 0.2
            t = cam.get_translation()

            if cam_names is None:
                cam_name = f"cam_{i+1}"
            else:
                cam_name = cam_names[i]
            data += vector_plot(
                [R_mat[:, 0], R_mat[:, 1], R_mat[:, 2]], orig=t, cam_name=cam_name
            )
    if skeleton_bp is not None and skeleton_lines is not None:
        for i, line in enumerate(skeleton_lines):
            from_bp = line[0]
            to_bp = line[1]

            # Standard case
            if isinstance(from_bp, str):
                vec = [
                    skeleton_bp[from_bp][i] - skeleton_bp[to_bp][i] for i in range(3)
                ]

                if line_colors is not None:
                    color = [line_colors[i]]
                else:
                    color = ["gray"]

                cam_name = to_bp if plot_names else ""
                # names = [from_bp]
                names = [""]

                vec_data = vector_plot(
                    [vec],
                    orig=skeleton_bp[to_bp],
                    cam_name=cam_name,
                    names=names,
                    colors=color,
                )
                data += vec_data

            # Want to draw to midpoint
            elif isinstance(from_bp, tuple):
                # calculate midpoint
                vec_from_from = [
                    skeleton_bp[from_bp[0]][i] - skeleton_bp[from_bp[1]][i]
                    for i in range(3)
                ]

                mid = [
                    (skeleton_bp[from_bp[1]][j] + vec_from_from[j] * from_bp[-1])
                    for j in range(3)
                ]

                vec = [mid[i] - skeleton_bp[to_bp][i] for i in range(3)]
                if line_colors is not None:
                    color = [line_colors[i]]
                else:
                    color = ["green"]
                vec_data = vector_plot(
                    [vec],
                    orig=skeleton_bp[to_bp],
                    cam_name=to_bp,
                    names=[""],
                    colors=color,
                )
                data += vec_data

    if scene_lims is not None:
        layout = go.Layout(
            margin=dict(l=0, r=0, b=0, t=0),
            scene=dict(
                xaxis=dict(nticks=4, range=scene_lims[0], showticklabels=False),
                yaxis=dict(nticks=4, range=scene_lims[1], showticklabels=False),
                zaxis=dict(nticks=4, range=scene_lims[2], showticklabels=False),
            ),
        )
    else:
        layout = go.Layout(margin=dict(l=0, r=0, b=0, t=0),)

    fig = go.FigureWidget(data=data, layout=layout)
    fig.update_traces(textfont_size=font_size)
    fig.update_layout(scene_aspectmode=scene_aspect)
    fig.update_layout(scene_camera=scene_camera)
    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))
    if not legend:
        fig.update_layout(showlegend=False)

    fig.update_layout(
        title={
            "text": title,
            "y": 0.9,
            "x": 0.5,
            "xanchor": "center",
            "yanchor": "top",
        },
    )

    if show_plot:
        fig.show()

    fig["layout"]["uirevision"] = "nothing"

    return fig

# Source: http://193.51.245.4/tutorials/convert_a_matplotlib_figure
def fig2data(fig):
    """
    @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
    @param fig a matplotlib figure
    @return a numpy 3D array of RGBA values
    """
    # draw the renderer
    fig.canvas.draw()

    # Get the RGBA buffer from the figure
    w, h = fig.canvas.get_width_height()
    buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8)
    buf.shape = (w, h, 4)

    # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
    buf = np.roll(buf, 3, axis=2)
    return buf

def fig2img(fig):
    """
    @brief Convert a Matplotlib figure to a PIL Image in RGBA format and return it
    @param fig a matplotlib figure
    @return a Python Imaging Library ( PIL ) image
    """
    # put the figure pixmap into a numpy array
    buf = fig2data(fig)
    w, h, d = buf.shape
    return Image.frombytes("RGBA", (w, h), buf.tostring())


import os
from pathlib import Path
from typing import Optional

import alembic.command
import alembic.config

# Unused netcdf4 import to suppress a warning from numpy/xarray/netcdf4
# https://github.com/pydata/xarray/issues/7259
import netCDF4  # type: ignore # noqa: F401 # This needs to be imported before numpy
import numpy as np
import pytest
import sqlalchemy
import xarray as xr
from sqlalchemy.orm import Session

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

test_db_user = os.environ.get("TEST_DB_USER", "postgres")
test_db_pass = os.environ.get("TEST_DB_PASS", "postgres")
test_db_host = os.environ.get("TEST_DB_HOST", "localhost:5432")

@pytest.fixture(scope="session")
def db_name():
    return "test_unified_graphics"

@pytest.fixture(scope="session")
def test_db(db_name):
    autocommit_engine = sqlalchemy.create_engine(
        f"postgresql+psycopg://{test_db_user}:{test_db_pass}@{test_db_host}/postgres",
        isolation_level="AUTOCOMMIT",
    )
    with autocommit_engine.connect() as conn:
        conn.execute(sqlalchemy.text(f"DROP DATABASE IF EXISTS {db_name}"))
        conn.execute(sqlalchemy.text(f"CREATE DATABASE {db_name}"))

    url = f"postgresql+psycopg://{test_db_user}:{test_db_pass}@{test_db_host}/{db_name}"

    config = alembic.config.Config("alembic.ini")
    config.set_main_option("sqlalchemy.url", url)
    alembic.command.upgrade(config, revision="head")

    yield url

    with autocommit_engine.connect() as conn:
        # Disconnect all users from the database we are dropping. I'm unsure why this is
        # necessary, as I think all fo the connections created for our tests should be
        # cleaned up by now.
        #
        # Copied from:
        # https://sqlalchemy-utils.readthedocs.io/en/latest/_modules/sqlalchemy_utils/functions/database.html#drop_database
        version = conn.dialect.server_version_info
        pid_column = "pid" if (version >= (9, 2)) else "procpid"
        conn.execute(
            sqlalchemy.text(
                f"""SELECT pg_terminate_backend(pg_stat_activity.{pid_column})
        FROM pg_stat_activity
        WHERE pg_stat_activity.datname = '{db_name}'
        AND {pid_column} <> pg_backend_pid();
        """
            )
        )
        conn.execute(sqlalchemy.text(f"DROP DATABASE {db_name}"))

    autocommit_engine.dispose()

@pytest.fixture(scope="session")
def engine(test_db):
    _engine = sqlalchemy.create_engine(test_db)
    yield _engine
    _engine.dispose()

@pytest.fixture(scope="class")
def session(engine):
    with Session(engine) as s:
        yield s
        s.rollback()

# FIXME: Replace diag_dataset with this fixture
@pytest.fixture(scope="class")
def test_dataset():
    def factory(
        *,
        model: str = "RTMA",
        system: str = "WCOSS",
        domain: str = "CONUS",
        background: str = "HRRR",
        frequency: str = "REALTIME",
        initialization_time: str = "2022-05-16T04:00",
        variable: str = "ps",
        loop: str = "ges",
        longitude: list[float] = [90, 91],
        latitude: list[float] = [22, 23],
        is_used: list[bool] = [True, False],
        observation: list[float] = [1, 0],
        forecast_unadjusted: list[float] = [0, 1],
        forecast_adjusted: Optional[list[float]] = None,
        **kwargs,
    ):
        assert len(observation) == len(forecast_unadjusted)
        if forecast_adjusted:
            assert len(forecast_adjusted) == len(forecast_unadjusted)

        dims = ["nobs", *kwargs.keys()]

        obs = np.array(observation, dtype=np.float64)
        unadj = np.array(forecast_unadjusted, dtype=np.float64)
        adj = np.array(forecast_adjusted or forecast_unadjusted, dtype=np.float64)

        return xr.Dataset(
            {
                "observation": (dims, obs),
                "forecast_unadjusted": (dims, unadj),
                "forecast_adjusted": (dims, adj),
                "obs_minus_forecast_unadjusted": (dims, obs - unadj),
                "obs_minus_forecast_adjusted": (dims, obs - adj),
            },
            coords=dict(
                longitude=(["nobs"], np.array(longitude, dtype=np.float64)),
                latitude=(["nobs"], np.array(latitude, dtype=np.float64)),
                is_used=(["nobs"], np.array(is_used)),
                **kwargs,
            ),
            attrs={
                "name": variable,
                "loop": loop,
                "initialization_time": initialization_time,
                "model": model,
                "system": system,
                "domain": domain,
                "frequency": frequency,
                "background": background,
            },
        )

    return factory

@pytest.fixture
def diag_parquet(tmp_path):
    def factory(
        ds: xr.Dataset,
    ) -> Path:
        parquet_file = (
            tmp_path
            / "_".join((ds.model, ds.background, ds.system, ds.domain, ds.frequency))
            / ds.name
        )
        df = ds.to_dataframe()
        df["loop"] = ds.loop
        df["initialization_time"] = ds.initialization_time

        df.to_parquet(
            parquet_file, partition_cols=["loop"], index=True, engine="pyarrow"
        )

        return parquet_file

    return factory


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

from itertools import product
import os

import numpy as np

PATH = os.environ.get("PSS_PATH", "..")

def pss_block(
    seed, k, case, i1, block_id, m=2_000_000, t=1_000, save=True, path="./"
):
    file_name = f"pss-k-{k}-case-{case}-i1-{i1}-block-{block_id}.npz"
    file_name = os.path.join(path, file_name)
    if save and os.path.exists(file_name):
        return
    rs = np.random.default_rng(seed)
    const = np.ones(t - 1)
    tau = np.arange(1, t).astype(float)
    f = np.empty(m)
    for j in range(m):
        u = rs.standard_normal((k + 1, t))
        y = np.cumsum(u[0])
        if i1:
            x = np.cumsum(u[1:], axis=1).T
        else:
            x = u[1:].T
        lhs = np.diff(y)
        rhv = [y[:-1], x[:-1]]
        if case == 2:
            rhv.append(const)
        elif case == 4:
            rhv.append(tau)
        if case >= 3:
            rhv.append(const)
        if case == 5:
            rhv.append(tau)
        rest = k + 1
        if case in (2, 4):
            rest += 1
        rhs = np.column_stack(rhv)
        b = np.linalg.lstsq(rhs, lhs, rcond=None)[0]
        u = lhs - rhs @ b
        s2 = u.T @ u / (u.shape[0] - rhs.shape[1])
        xpx = rhs.T @ rhs
        vcv = np.linalg.inv(xpx) * s2
        r = np.eye(rest, rhs.shape[1])
        rvcvr = r @ vcv @ r.T
        rb = r @ b
        f[j] = rb.T @ np.linalg.inv(rvcvr) @ rb / rest
    percentiles = [0.05]
    percentiles += [i / 10 for i in range(1, 10)]
    percentiles += [1 + i / 2 for i in range(18)]
    percentiles += [i for i in range(10, 51)]
    percentiles += [100 - v for v in percentiles]
    percentiles = sorted(set(percentiles))
    percentiles = np.asarray(percentiles)
    q = np.percentile(f, percentiles)
    if save:
        np.savez(file_name, q=q, percentiles=percentiles)
    return q

seed = [
    3957597042,
    2709280948,
    499296859,
    1555610991,
    2390531900,
    2160388094,
    4098495866,
    47221919,
]
ss = np.random.SeedSequence(seed)
k = list(range(1, 11))
case = list(range(1, 6))
i1 = [True, False]
block_id = list(range(32))
params = list(product(k, case, i1, block_id))
seeds = ss.generate_state(8 * len(params)).reshape((-1, 8)).tolist()
configs = []
for _s, (_k, _case, _i1, _block_id) in zip(seeds, params):
    configs.append(
        {
            "seed": _s,
            "k": _k,
            "case": _case,
            "i1": _i1,
            "block_id": _block_id,
            "path": PATH,
        }
    )

if __name__ == "__main__":
    from joblib import Parallel, delayed

    Parallel(n_jobs=10)(delayed(pss_block)(**c) for c in configs)


import argparse
import os
import shutil
import subprocess
import sys
import tempfile
import zipfile

import utils

from android_build_client import AndroidBuildClient

class Benchmark(object):

    def __init__(self, name):
        self.name = name

    def apct_test_tag(self):
        raise NotImplementedError()

    def profdata_file(self, suffix=''):
        profdata = os.path.join(self.name, '{}.profdata'.format(self.name))
        if suffix:
            profdata += '.' + suffix
        return profdata

    def profraw_files(self):
        raise NotImplementedError()

    def merge_profraws(self, profile_dir, output):
        profraws = [
            os.path.join(profile_dir, p)
            for p in self.profraw_files(profile_dir)
        ]
        utils.run_llvm_profdata(profraws, output)

class NativeExeBenchmark(Benchmark):

    def apct_test_tag(self):
        return 'apct/perf/pgo/profile-collector'

    def profraw_files(self, profile_dir):
        if self.name == 'hwui':
            return [
                'hwuimacro.profraw', 'hwuimacro_64.profraw',
                'hwuimicro.profraw', 'hwuimicro_64.profraw',
                'skia_nanobench.profraw', 'skia_nanobench_64.profraw'
            ]
        elif self.name == 'hwbinder':
            return ['hwbinder.profraw', 'hwbinder_64.profraw']

class APKBenchmark(Benchmark):

    def apct_test_tag(self):
        return 'apct/perf/pgo/apk-profile-collector'

    def profdata_file(self, suffix=''):
        profdata = os.path.join('art',
                                '{}_arm_arm64.profdata'.format(self.name))
        if suffix:
            profdata += '.' + suffix
        return profdata

    def profraw_files(self, profile_dir):
        return os.listdir(profile_dir)

def BenchmarkFactory(benchmark_name):
    if benchmark_name == 'dex2oat':
        return APKBenchmark(benchmark_name)
    elif benchmark_name in ['hwui', 'hwbinder']:
        return NativeExeBenchmark(benchmark_name)
    else:
        raise RuntimeError('Unknown benchmark ' + benchmark_name)

def extract_profiles(build, test_tag, build_client, output_dir):
    pgo_zip = build_client.download_pgo_zip(build, test_tag, output_dir)

    zipfile_name = os.path.join(pgo_zip)
    zip_ref = zipfile.ZipFile(zipfile_name)
    zip_ref.extractall(output_dir)
    zip_ref.close()

KNOWN_BENCHMARKS = ['ALL', 'dex2oat', 'hwui', 'hwbinder']

def parse_args():
    """Parses and returns command line arguments."""
    parser = argparse.ArgumentParser()

    parser.add_argument(
        'build',
        metavar='BUILD',
        help='Build number to pull from the build server.')

    parser.add_argument(
        '-b', '--bug', type=int, help='Bug to reference in commit message.')

    parser.add_argument(
        '--use-current-branch',
        action='store_true',
        help='Do not repo start a new branch for the update.')

    parser.add_argument(
        '--add-do-not-merge',
        action='store_true',
        help='Add \'DO NOT MERGE\' to the commit message.')

    parser.add_argument(
        '--profdata-suffix',
        type=str,
        default='',
        help='Suffix to append to merged profdata file')

    parser.add_argument(
        'benchmark',
        metavar='BENCHMARK',
        help='Update profiles for BENCHMARK.  Choices are {}'.format(
            KNOWN_BENCHMARKS))

    parser.add_argument(
        '--skip-cleanup',
        '-sc',
        action='store_true',
        default=False,
        help='Skip the cleanup, and leave intermediate files (in /tmp/pgo-profiles-*)'
    )

    return parser.parse_args()

def get_current_profile(benchmark):
    profile = benchmark.profdata_file()
    dirname, basename = os.path.split(profile)

    old_profiles = [f for f in os.listdir(dirname) if f.startswith(basename)]
    if len(old_profiles) == 0:
        return ''
    return os.path.join(dirname, old_profiles[0])

def main():
    args = parse_args()
    utils.check_gcertstatus()

    if args.benchmark == 'ALL':
        worklist = KNOWN_BENCHMARKS[1:]
    else:
        worklist = [args.benchmark]

    profiles_project = os.path.join(utils.android_build_top(), 'toolchain',
                                    'pgo-profiles')
    os.chdir(profiles_project)

    if not args.use_current_branch:
        branch_name = 'update-profiles-' + args.build
        utils.check_call(['repo', 'start', branch_name, '.'])

    build_client = AndroidBuildClient()

    for benchmark_name in worklist:
        benchmark = BenchmarkFactory(benchmark_name)

        # Existing profile file, which gets 'rm'-ed from 'git' down below.
        current_profile = get_current_profile(benchmark)

        # Extract profiles to a temporary directory.  After extraction, we
        # expect to find one subdirectory with profraw files under the temporary
        # directory.
        extract_dir = tempfile.mkdtemp(prefix='pgo-profiles-' + benchmark_name)
        extract_profiles(args.build, benchmark.apct_test_tag(), build_client,
                         extract_dir)

        extract_subdirs = [
            os.path.join(extract_dir, sub)
            for sub in os.listdir(extract_dir)
            if os.path.isdir(os.path.join(extract_dir, sub))
        ]
        if len(extract_subdirs) != 1:
            raise RuntimeError(
                'Expected one subdir under {}'.format(extract_dir))

        # Merge profiles.
        profdata = benchmark.profdata_file(args.profdata_suffix)
        benchmark.merge_profraws(extract_subdirs[0], profdata)

        # Construct 'git' commit message.
        message_lines = [
            'Update PGO profiles for {}'.format(benchmark_name), '',
            'The profiles are from build {}.'.format(args.build), ''
        ]

        if args.add_do_not_merge:
            message_lines[0] = '[DO NOT MERGE] ' + message_lines[0]

        if args.bug:
            message_lines.append('')
            message_lines.append('Bug: http://b/{}'.format(args.bug))
        message_lines.append('Test: Build (TH)')
        message = '\n'.join(message_lines)

        # Invoke git: Delete current profile, add new profile and commit these
        # changes.
        if current_profile:
            utils.check_call(['git', 'rm', current_profile])
        utils.check_call(['git', 'add', profdata])
        utils.check_call(['git', 'commit', '-m', message])

        if not args.skip_cleanup:
            shutil.rmtree(extract_dir)

if __name__ == '__main__':
    main()


from sklearn.model_selection import train_test_split
from sklearn.model_selection import KFold
import numpy as np

x = np.arange(20).reshape((10, 2))
x2 = np.arange(10).reshape((5, 2))
y = np.arange(10)
y2 = np.arange(5)

cv = KFold(n_splits=5, shuffle=True, random_state=0)
for (train1_idx, eval1_idx), (train2_idx, eval2_idx) in zip(cv.split(x), cv.split(x2)):
    print('train idx:', train1_idx, train2_idx)
    print('eval idx:', eval1_idx, eval2_idx)
    print('.' * 50)

    print(x[train1_idx])


# -*- coding: utf-8 -*-
# @Time    : 2020/3/6 19:01
# @Author  : Deng Wenxing
# @Email   : dengwenxingae86@163.com
# @File    : data_count.py
# @Software: PyCharm
# @content : 统计每个表的数量
from pyspark import SparkConf
from pyspark.sql import SparkSession
from pyspark.sql.functions import *
import time
import os,sys
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s: %(message)s')
logger = logging.getLogger(__name__)

reload(sys)
sys.setdefaultencoding('utf-8')

conf=SparkConf().set('spark.driver.maxResultSize', '30g')
conf.set('spark.yarn.am.cores', 4)
conf.set('spark.executor.memory', '15g')
conf.set('spark.executor.instances', 100)
conf.set('spark.executor.cores', 4)
conf.set('spark.executor.extraJavaOptions', '-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -XX:+UseG1GC')

spark = SparkSession \
    .builder \
    .config(conf=conf) \
    .enableHiveSupport() \
    .getOrCreate()

from common import *
from jg_info import vertex_table_info,edge_table_info

path_prefix = '/phoebus/_fileservice/users/slmp/shulianmingpin/midfile/open_phone'
save_root = 'relation_theme_extenddir'
zp_source_root = 'wa_data/zp_jg'
steal_oil_source_root = 'wa_data/touyou_jg'

exists_files = vertex_table_info.keys() + edge_table_info.keys()
zp_file_cols = {
    'vertex_phonenumber':['phone', 'xm','province', 'city'],
    'vertex_person':['zjhm', 'zjlx', 'gj', 'xm', 'ywxm', 'zym', 'xb', 'csrq', 'mz', 'jg',
                                      'whcd', 'hyzk', 'zzmm', 'hkszdxz', 'sjjzxz'],
    'vertex_case':['asjbh','ajmc','0 asjfskssj','0 asjfsjssj','"" asjly','ajlb','"" fxasjsj','fxasjdd_dzmc',
                                        'jyaq','"" ajbs','larq','city','link_app'],
    'vertex_qq':['qq'],
    'vertex_alipay':['alipay'],
    'vertex_qq_group':['groupid','groupname'],
    'vertex_bankcard':['yhkh banknum','khh','sfzh','khrq'],
    'edge_qqgroup_link_banknum':['banknum','groupid','relate_time','start_time','end_time','relate_num'],
    'edge_qq_link_group':['qq','groupid','start_time','end_time'],
    'edge_qq_link_qq':['qq1','qq2','start_time','end_time'],
    'edge_person_link_phone':['phone','sfzh','start_time','end_time'],
    'edge_case_link_wechat':['wechatid wechat','ajbh asjbh','start_time','end_time'],
    'edge_case_link_qq':['qq','ajbh asjbh','start_time','end_time'],
    'edge_case_link_phone':['phone','ajbh asjbh','start_time','end_time'],
    'edge_case_link_person':['sfzh','ajbh asjbh','start_time','end_time'],
    'edge_case_link_banknum':['banknum','ajbh asjbh','start_time','end_time'],
    'edge_case_link_alipay':['alipay','ajbh asjbh','start_time','end_time'],
}

steal_file_cols = {
    'vertex_case':['asjbh','ajmc','asjfskssj','asjfsjssj','asjly','ajlb','fxasjsj','fxasjdd_dzmc',
                                        'jyaq','ajbs','larq','"" city','"" link_app'],
    'edge_case_link_person':['sfzh','ajbh asjbh','start_time','end_time'],
    'edge_case_link_phone':['phone','ajbh asjbh','start_time','end_time'],

}

def deal_zp_data():
    for file,cols in zp_file_cols.items():
        logger.info('dealing %s'%file)
        logger.info(add_save_path(file+"_jg",cp='*',root=zp_source_root))
        df = read_orc(spark,add_save_path(file+"_jg",cp='*',root=zp_source_root)).selectExpr(*cols)
        file_name = file+"_zp" if file in exists_files else file
        write_orc(df,add_save_path(file_name,root=save_root))

def deal_steal_data():
    for file,cols in steal_file_cols.items():
        df = read_orc(spark,add_save_path(file+"_jg",cp='*',root=steal_oil_source_root)).selectExpr(*cols)
        file_name = file+"_steal" if file in exists_files else file
        write_orc(df,add_save_path(file_name,root=save_root))

if __name__ == "__main__":
    logger.info('========================start time:%s==========================' % (
        time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())))

    # deal_zp_data()
    deal_steal_data()

    logger.info('========================end time:%s==========================' % (
        time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())))


"""
Programming_1 finall assignment
Author: Stylianos Mavrianos
Date: 25/01/2022
Version: 2

Supporting file for finall assignment
"""

#Importing Libralies
import pandas as pd
import numpy as np
from bokeh.io import output_notebook
from bokeh.plotting import figure, show
from bokeh.transform import jitter
from bokeh.models import ColumnDataSource,Range1d
from bokeh.layouts import row
import matplotlib.pyplot as plt

def get_NP_level(df, np):
     #######################################################
    #
    # Arguments:
    #
    # df                  pd.dataframe
    # np                  level of nanoplastic concentration
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################

    df_np_level = df[df['NP_treatment'] == np]

    return df_np_level

def get_INF_level(df, inf):
     #######################################################
    #
    # Arguments:
    #
    # df                  pd.dataframe
    # np                  level of infection
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################
    df_inf_level = df[df['Inf_treatment'] == inf]

    return df_inf_level

def get_values(df, var):
     #######################################################
    #
    # Arguments:
    #
    # df                  pd.dataframe
    # var                 variable of interest
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################
    if var == 'Age_death':
        x = 3
    if var == 'Total_juv':
        x = 4
    if var == 'Spore_yield':
        x = 8
    values = np.array(df.iloc[0:, x])

    return values

def get_means(df, var):
    #######################################################
    #
    # Arguments:
    #
    # df                  pd.dataframe
    # var                 variable of interest
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################
    means = pd.DataFrame(data=df.groupby(['Inf_treatment','NP_treatment'])[var].mean())
    means = means.reset_index()
    means = means.drop(columns=['Inf_treatment'])

    return means

def get_std(df, var):
    #######################################################
    #
    # Arguments:
    #
    # df                  pd.dataframe
    # var                 variable of interest
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################
    high, low, zero = df.groupby('NP_treatment')

    std_zero = np.std(zero[1][var], ddof=1) / np.sqrt(np.size(zero[1][var]))
    std_low = np.std(low[1][var], ddof=1) / np.sqrt(np.size(low[1][var]))
    std_high = np.std(high[1][var], ddof=1) / np.sqrt(np.size(high[1][var]))

    d = {'NP_treatment': ['HIGH','LOW','ZERO'], 'std': [std_high,std_low,  std_zero,]}
    std =  pd.DataFrame(data=d)

    return std

def errorbar_df(df, var):
    #######################################################
    #
    # Arguments:
    #
    # df                  pd.dataframe
    # var                 variable of interest
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################
    mean = get_means(df, var)
    std = get_std(df, var)
    mean['std_lower'] = mean[var] - std['std']
    mean['std_upper'] = mean[var] + std['std']

    return mean

def make_plot(df_main, df_errorbar, var, title,axis_label, num):
    #######################################################
    #
    # Arguments:
    #
    # df_main                  pd.dataframe
    # df_errorbar              pd.dataframe
    # var                      variable of interest
    # title                    title of the plot
    # axis_label               label of the plot
    # num                      y_max_range
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################
    np_levels = ['ZERO', 'LOW', 'HIGH']

    source = ColumnDataSource(df_main)
    source_1 = ColumnDataSource(df_errorbar)

    p = figure(tools="pan, hover")
    p = figure(x_range= np_levels,y_range = (0,num),
               toolbar_location=None,
               title=title,
               plot_height=300, plot_width=300)
    p.yaxis.axis_label = axis_label

    #Scatter data
    p.dot(x= jitter('NP_treatment',width=0.2, range=p.x_range), y=jitter(var, width=1, range=p.y_range), source= source , alpha=0.3, size = 35)

    p.line(x = 'NP_treatment', y = var, source= source_1, color = 'black')

    # Error asterisks
    p.asterisk(x = 'NP_treatment', y = 'std_lower', source= source_1, color = 'black', size = 3)
    p.asterisk(x = 'NP_treatment', y = 'std_upper', source= source_1, color = 'black', size = 3)

    # Mean
    p.dot(x = 'NP_treatment', y = var, source= source_1, size = 35, color = 'red')

    # p.ygrid.grid_line_color = None

    return p

def make_barplot(df_main, df_errorbar, var, title, axis_label):
    #######################################################
    #
    # Arguments:
    #
    # df_main                  pd.dataframe
    # df_errorbar              pd.dataframe
    # var                 variable of interest
    # title                    title of the plot
    # axis_label               label of the plot
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################
    from bokeh.palettes import Spectral3

    h, l, z = df_main.groupby(['NP_treatment'])[var].sum() / df_main.groupby(['NP_treatment'])[var].count()

    NP_treatment = ['ZERO', 'LOW', 'HIGH']
    counts = [z, l, h]

    source = ColumnDataSource(data=dict(NP_treatment = NP_treatment, counts = counts, color=Spectral3))
    source_1 = ColumnDataSource(df_errorbar)

    p = figure(x_range = NP_treatment, y_range = (0,1.05), height = 250, title = title,
               toolbar_location = None, tools = "")

    p.vbar(x='NP_treatment', top = 'counts', width=0.2, color='color', source = source)
    p.xgrid.grid_line_color = None
    p.yaxis.axis_label = axis_label

    # Error asterisks
    p.asterisk(x = 'NP_treatment', y = 'std_lower', source= source_1, color = 'black', size = 3)
    p.asterisk(x = 'NP_treatment', y = 'std_upper', source= source_1, color = 'black', size = 3)

    return p

def anova(df, Response_var, *Explanatory_var ):
    #######################################################
    #
    # Arguments:
    #
    # df                  pd.dataframe
    # Response_var        Response variable
    # Explanatory_var     Explanatory variables
    #
    # Author:            Stylianos Mavrianos
    # Date:              2021-11-06
    #
    ########################################################

    import pingouin as pg

    aov = pg.anova(dv=Response_var, between=[*Explanatory_var], data= df,
               detailed=True)

    return aov.round(4)

def figures():

    zero_counts = np.array([11, 14])
    zero_labels = ["Infected", "Not infected",]
    zero_colors = ['r', 'g']
    plt.pie(zero_counts, labels = zero_labels, colors = zero_colors, shadow = True)
    plt.title("Zero concentration")
    plt.show()

    low_counts = np.array([1, 2, 6, 16])
    low_labels = ['excluded', 'early death', "Infected", "Not infected"]
    low_colors = ['k', 'gray', 'r', 'g']
    plt.title("Low concentration")
    plt.pie(low_counts, labels = low_labels, colors = low_colors, shadow = True, counterclock = True)
    plt.show()

    high_counts = np.array([2, 6, 3, 14])
    high_labels = ['excluded', 'early death', "Infected", "Not infected"]
    high_colors = ['k', 'gray', 'r', 'g']
    plt.title("High concentration")
    plt.pie(high_counts, labels = high_labels, colors = high_colors, shadow = True)
    plt.show()

def Q_Q_Plot(y, est = 'robust', **kwargs):

    ################################################################################
    #
    # Arguments:
    #
    # y                  data array
    # est                Estimation method for normal parameters mu and sigma:
    #                    either 'robust' (default), or 'ML' (Maximum Likelihood),
    #                    or 'preset' (given values)
    # If est='preset' than the optional parameters mu, sigma must be provided
    #
    # Author:            M.E.F. Apol
    # Date:              2020-01-06
    #
    ################################################################################

    import numpy as np
    from scipy.stats import iqr # iqr is the Interquartile Range function
    import matplotlib.pyplot as plt

    # First, get the optional arguments mu and sigma:
    mu_0 = kwargs.get('mu', None)
    sigma_0 = kwargs.get('sigma', None)

    n = len(y)

    # Calculate order statistic:
    y_os = np.sort(y)

    # Estimates of mu and sigma:
    # ML estimates:
    mu_ML = np.mean(y)
    sigma2_ML = np.mean((y - mu_ML)**2)
    sigma_ML = np.sqrt(sigma2_ML) # biased estimate
    s2 = n/(n-1) * sigma2_ML
    s = np.sqrt(s2) # unbiased estimate
    # Robust estimates:
    mu_R = np.median(y)
    sigma_R = iqr(y)/1.349

    # Assign values of mu and sigma for z-transform:
    if est == 'ML':
        mu, sigma = mu_ML, s
    elif est == 'robust':
        mu, sigma = mu_R, sigma_R
    elif est == 'preset':
        mu, sigma = mu_0, sigma_0
    else:
        print('Wrong estimation method chosen!')

    print('Estimation method: ' + est)
    print('mu = ',mu,', sigma = ',sigma)

    # Perform z-transform: sample quantiles z.i
    z_i = (y_os - mu)/sigma

    # Calculate cumulative probabilities p.i:
    i = np.array(range(n)) + 1
    p_i = (i - 0.5)/n

    # Calculate theoretical quantiles z.(i):
    from scipy.stats import norm
    z_th = norm.ppf(p_i, 0, 1)

    # Calculate SE or theoretical quantiles:
    SE_z_th = (1/norm.pdf(z_th, 0, 1)) * np.sqrt((p_i * (1 - p_i)) / n)

    # Calculate 95% CI of diagonal line:
    CI_upper = z_th + 1.96 * SE_z_th
    CI_lower = z_th - 1.96 * SE_z_th

    # Make Q-Q plot:
    plt.plot(z_th, z_i, 'o', color='k', label='experimental data')
    plt.plot(z_th, z_th, '--', color='r', label='normal line')
    plt.plot(z_th, CI_upper, '--', color='b', label='95% CI')
    plt.plot(z_th, CI_lower, '--', color='b')
    plt.xlabel('Theoretical quantiles, $z_{(i)}$')
    plt.ylabel('Sample quantiles, $z_i$')
    plt.title('Q-Q plot (' + est + ')')
    plt.legend(loc='best')
    plt.show()

def Q_Q_Hist(y, est='robust', **kwargs):

    ################################################################################
    #
    # Arguments:
    #
    # y                  data array
    # est                Estimation method for normal parameters mu and sigma:
    #                    either 'robust' (default), or 'ML' (Maximum Likelihood),
    #                    or 'preset' (given values)
    # If est='preset' than the optional parameters mu, sigma must be provided
    #
    # Author:            M.E.F. Apol
    # Date:              2020-01-06
    #
    ################################################################################

    import numpy as np
    from scipy.stats import iqr # iqr is the Interquartile Range function
    from scipy.stats import norm
    import matplotlib.pyplot as plt

    # First, get the optional arguments mu and sigma:
    mu_0 = kwargs.get('mu', None)
    sigma_0 = kwargs.get('sigma', None)

    n = len(y)

    # Estimates of mu and sigma:
    # ML estimates:
    mu_ML = np.mean(y)
    sigma2_ML = np.mean((y - mu_ML)**2)
    sigma_ML = np.sqrt(sigma2_ML) # biased estimate
    s2 = n/(n-1) * sigma2_ML
    s = np.sqrt(s2) # unbiased estimate
    # Robust estimates:
    mu_R = np.median(y)
    sigma_R = iqr(y)/1.349

    # Assign values of mu and sigma for z-transform:
    if est == 'ML':
        mu, sigma = mu_ML, s
    elif est == 'robust':
        mu, sigma = mu_R, sigma_R
    elif est == 'preset':
        mu, sigma = mu_0, sigma_0
    else:
        print('Wrong estimation method chosen!')
    print('Estimation method: ' + est)
    print('mu = ',mu,', sigma = ',sigma)

    # Calculate the CLT normal distribution:
    x = np.linspace(np.min(y), np.max(y), 501)
    rv = np.array([norm.pdf(xi, loc = mu, scale = sigma) for xi in x])

    # Make a histogram with corresponding normal distribution:
    nn, bins, patches = plt.hist(x=y, density=True,
                                 bins='auto',
                                 color='darkgrey',alpha=1, rwidth=1, label='experimental')
    h = plt.plot(x, rv, 'r', label='normal approximation')
    plt.grid(axis='y', alpha=0.5)
    plt.xlabel('Values, $y$')
    plt.ylabel('Probability $f(y)$')
    plt.title('Histogram with corresponding normal distribution (' + est + ')')
    plt.legend(loc='best')
    plt.show()

if __name__ == "__main__":
    main()


from future import standard_library
standard_library.install_aliases()
import logging
import sys

from desktop.lib.conf import Config, coerce_bool
from desktop.conf import default_ssl_validate
from libzookeeper.conf import ENSEMBLE

if sys.version_info[0] > 2:
  from urllib.parse import urlparse
  new_str = str
  from django.utils.translation import gettext_lazy as _t
else:
  from django.utils.translation import ugettext_lazy as _t
  from urlparse import urlparse

LOG = logging.getLogger()

SSL_CERT_CA_VERIFY = Config(
  key="ssl_cert_ca_verify",
  help=_t("In secure mode (HTTPS), if Solr SSL certificates have to be verified against certificate authority"),
  dynamic_default=default_ssl_validate,
  type=coerce_bool
)

def zkensemble_path():
  """
  Try to guess Solr path in ZooKeeper.
  """
  try:
    parsed = urlparse(ENSEMBLE.get())
    if parsed.port == 9983: # Standalone Solr cloud
      return ''
  except:
    LOG.warning('Failed to get Zookeeper ensemble path')

  return '/solr'

SOLR_ZK_PATH = Config(
  key="solr_zk_path",
  help=_t("Default path to Solr in ZooKeeper"),
  dynamic_default=zkensemble_path,
  type=str
)


from gmusicapi import Mobileclient
from Initialize import sqliteread, sqlitewrite, openSocket, sendMessage,  createsongqueue, reconnect
from pytube import YouTube
import validators
import vlc
from Settings import *
import time
import urllib.request, urllib.parse, urllib.error
from shutil import copyfile
import sqlite3
from sqlite3 import Error
from importlib import reload
from unidecode import unidecode
import sys
reload(sys)

s = openSocket()
api = Mobileclient()
api.oauth_login(device_id=Mobileclient.FROM_MAC_ADDRESS, oauth_credentials="oauth.txt")
if Mobileclient.is_authenticated(api) == True:
    print("Logged into GPM successfully")
else:
    sendMessage(s, "Can't log into Google Play Music! Please check the console and follow the instructions!")
    api.perform_oauth(storage_filepath="oauth.txt")
    api.oauth_login(device_id=Mobileclient.FROM_MAC_ADDRESS, oauth_credentials="oauth.txt")

stream_url = ""

def writenowplaying(isPlaying, song_name):
    with open("Output/nowplaying.txt", "w") as f:
        if isPlaying == True:
            f.write(song_name)
        else:
            f.truncate()

def getytkey(url):
    if "youtube" in url.lower():
        return(url.split("=")[1][:11])
    elif "youtu.be" in url.lower():
        return(url.split("e/")[1][:11])
    else:
        return None

def songtitlefilter(song_name, redo):
    blacklist = BLACKLISTED_SONG_TITLE_CONTENTS[:]
    try:
        results = (Mobileclient.search(api, song_name, SONGBLSIZE)['song_hits'])[:SONGBLSIZE]
    except Exception as e:
        print(e)
        reconnect()
    songs = []

    for item in results:
        songs.append(item['track'])

    #Remove things from the blacklist if theyre explicitly requested
    for term in blacklist:
        if term.lower() in song_name.lower():
            blacklist.remove(term)

    #Iterate through the blacklisted contents, then the songs. Last song standing wins.
    for term in blacklist:
        if len(songs) == 1:
            break
        for song in reversed(songs):
            if len(songs) == 1:
                break
            if term.lower() in song['title'].lower():
                print((">Removed: " + song['title']))
                songs.remove(song)

    for item in songs:
        print((">>>Allowed: " + item['title']))
    print((">>>>>>Playing: " + songs[0]['title']))
    return songs[redo]

def sr_geturl(songkey):
    try:
        stream_url = Mobileclient.get_stream_url(api, songkey)
        return(stream_url)
    except Exception as e:
        print(e)
        reconnect()

def saveAlbumArt(songkey):
    if songkey[0] == "T": #If the key is from GPM - all GPM keys start with T.
        songinfo = Mobileclient.get_track_info(api, songkey)
        imgLink = songinfo['albumArtRef'][0]['url']
        f = open('Output/albumart.jpg', 'wb')
        f.write(urllib.request.urlopen(imgLink).read())
        f.close()
    else: #Otherwise just use the generic image.
        copyfile('generic_art.jpg', 'Output/albumart.jpg')

def removetopqueue():
    row = sqliteread('''SELECT id, name, song, key FROM songs ORDER BY id ASC''') #Pick the top song
    if row[1] == "BotPlaylist": #>>>>>> DO THIS STUFF IF THE LIST IS EMPTY!
        sqlitewrite(('''DELETE FROM songs WHERE id={0}''').format(int(row[0]))) #Delete the top song
        return
    return False

def getnewentry():
    result = sqliteread('SELECT id FROM songs ORDER BY id DESC LIMIT 1')
    return(str(result[0]))

def playfromplaylist():
    row = sqliteread('''SELECT id, song, key FROM playlist ORDER BY id ASC''') #Pick the top song
    if row == None: #>>>>>> DO THIS STUFF IF THE LIST IS EMPTY!
        return
    try:
        songtitle = row[1]
        songkey = row[2]
    except:
        return

    #Delete the top result
    row = sqliteread('SELECT id FROM playlist ORDER BY id ASC LIMIT 1')
    sqlitewrite(('''DELETE FROM playlist WHERE id={0}''').format(int(row[0]))) #Delete the top song
    sqlitewrite(('''INSERT INTO playlist(song, key) VALUES("{song_name}", "{key}");''').format(song_name=songtitle, key=songkey))
    sqlitewrite(('''INSERT INTO songs(name, song, key) VALUES("BotPlaylist", "{song_name}", "{key}");''').format(song_name=songtitle, key=songkey))

class SRcontrol:
    def __init__(self):
        self.msg = None
        self.songtitle = ""

    def playsong(self):
        global skiprequests, skipusers
        row = sqliteread('''SELECT id, name, song, key FROM songs ORDER BY id ASC''') # Pick the top song
        try:
            self.songtitle = row[2]
            songkey = row[3]
        except Exception as e:
            time.sleep(0.4)
            return
        sqlitewrite(('''DELETE FROM songs WHERE id={0}''').format(row[0]))  # Delete the top song

        if validators.url(self.songtitle) == True: # TEST IF THE REQUEST IS A LINK
            if "youtu" in self.songtitle:  # IS IT A YOUTUBE LINK?
                try:
                    yt = YouTube(self.songtitle)
                    playurl = yt.streams.filter(only_audio=True).order_by('abr').first().url
                    self.songtitle = yt.title
                except Exception as e:
                    print("Pytube Error:")
                    print(e)
                writenowplaying(True, self.songtitle)
            else:  # If not YT link, play normal music file.
                playurl = self.songtitle
                self.songtitle = "Online Music File"
                writenowplaying(True, self.songtitle)
        else:  # Otherwise it's GPM
            playurl = sr_geturl(songkey)
            writenowplaying(True, self.songtitle)
        self.instance = vlc.Instance()
        self.p = self.instance.media_player_new(playurl)  # Play the music
        self.p.play()
        saveAlbumArt(songkey)
        writenowplaying(True, self.songtitle)
        createsongqueue()
        skiprequests = 0
        skipusers = []
        return True

    def songover(self):
        print("Song is over!")
        time.sleep(DELAY_BETWEEN_SONGS)
        self.p.stop()
        return False

    def gettime(self):
        return int(self.p.get_time())

    #-------------------------SONG REQUEST COMMANDS----------------------------------

    def volume(self, vol, user):
        try:
            if vol == None:
                return "Current volume: " + str(self.p.audio_get_volume())
            if (vol > 100) or (vol < 0):
                return "Invalid volume level. Must be between 0-100."
            self.p.audio_set_volume(vol)
            return "Music volume set to: " + str(vol)
        except AttributeError:
            return "Music needs to be playing before adjusting the volume."

    def volumeup(self, vol, user):
        try:
            if not vol:
                vol = VOL_INCREMENT
            currentvolume = self.p.audio_get_volume()
            if (currentvolume + vol) > 100:
                self.p.audio_set_volume(100)
                print ("Raised the volume to: 100")
                return "Raised the volume to: 100"
            self.p.audio_set_volume(currentvolume + vol)
            self.msg = "Raised the volume to: " + str(currentvolume + vol)
            print((self.msg))
            return self.msg
        except AttributeError:
            return "Music needs to be playing before adjusting the volume."

    def volumedown(self, vol, user):
        try:
            if not vol:
                vol = VOL_INCREMENT
            currentvolume = self.p.audio_get_volume()
            if (currentvolume - vol) < 0:
                self.p.audio_set_volume(0)
                print ("Lowered the volume to: 0")
                return "Lowered the volume to: 0"
            self.p.audio_set_volume(currentvolume - vol)
            self.msg = "Lowered the volume to: " + str(currentvolume - vol)
            print((self.msg))
            return self.msg
        except AttributeError:
            return "Music needs to be playing before adjusting the volume."

    def play(self):
        print ("Resumed the music")
        writenowplaying(True, self.songtitle)
        self.p.set_pause(False)

    def pause(self):
        print ("Paused the music")
        writenowplaying(False, "")
        self.p.set_pause(True)

class SRcommands:
    def __init__(self):
        self.db = None
        self.video = None
        global skiprequests, skipusers
        skiprequests = 0
        skipusers = []

    '''--------------------SONG REQUEST--------------------'''

    def songrequest(self, request, user):
        if request == "\r":  # Send a message if they request nothing
            return user + " >> " + DEFAULT_SR_MSG
        # Check for songs already in the queue
        if sqliteread('''SELECT count(*) FROM songs WHERE name="{0}"'''.format(user))[0] > (MAX_REQUESTS_USER - 1):
            return user + " >> You already have " + str(MAX_REQUESTS_USER) + " songs in the queue."

        request = request.replace("\r", "")[1:]
        # LINKS
        # DETERMINE LINK TYPE
        if validators.url(request):
            if "youtu" in request:
                key = getytkey(request)
                if not key: return "Something is wrong with your Youtube link."

                # Check the queue to see if the song is already in there.
                self.db = sqliteread('''SELECT id, count(*) FROM songs WHERE key="{0}"'''.format(key))

                if self.db[1] > (MAX_DUPLICATE_SONGS - 1):
                    return user + " >> That song is already in the queue."

                try:
                    self.video = YouTube(request)
                except:
                    return user + " >> That video is unavailable, it's probably age restricted."

                songtime = self.getsongtime(request, key)
                if songtime > (MAXTIME * 60000):
                    return user + " >> That song exceeds the maximum length of " + str(MAXTIME) + " minutes."

                title = self.video.title
                sqlitewrite('''INSERT INTO songs(name, song, key, time) VALUES("{user}", "{request}", "{key}", "{time}");'''.format(user=user, request=(request.replace('"', "'")), key=(key.replace('"', "'")), time=songtime))
                removetopqueue()
                return user + " >> Added: " + title + " to the queue (YT). ID: " + getnewentry()
            else:  # OTHER MP3 REQUESTS <<<<<<<
                if not MEDIA_FILE_ENABLE:
                    return user + " >> Online Media Links are disabled by the streamer, you'll need to request something else."

                songtime = self.getsongtime("Online", request)
                if songtime > (MAXTIME * 60000):
                    return user + " >> That song exceeds the maximum length of " + str(MAXTIME) + " minutes."

                sqlitewrite('''INSERT INTO songs(name, song, key, time) VALUES("{user}", "{request}", "{request}", "{time}");'''.format(user=user, request=(request.replace('"', "'")), time=songtime))
                removetopqueue()

                return user + " >> Added that link to the queue. ID: " + getnewentry()

        elif GPM_ENABLE:  # GOOGLE PLAY MUSIC STUFF
            try:
                top_song_result = songtitlefilter(request, 0)
                key = top_song_result['storeId']
                songtitle = str(top_song_result['artist'] + " - " + top_song_result['title'])
            # If theres an error (its unable to find the song) then announce it, otherwise write the song data to the db
            except IndexError:
                if YT_IF_NO_RESULT:
                    return self.youtubesr(request, user)
                else:
                    return user + " >> No results found for that song. Please try a different one."
            else:
                # Test if the song is already in the queue
                self.db = sqliteread('''SELECT id, count(*) FROM songs WHERE key="{0}"'''.format(key))
                if self.db[1] > (MAX_DUPLICATE_SONGS - 1):
                    return user + " >> That song is already in the queue."
                songtime = self.getsongtime(None, key)

                if songtime > (MAXTIME * 60000):
                    return user + " >> That song exceeds the maximum length of " + str(MAXTIME) + " minutes."

                # Add song to the queue
                sqlitewrite('''INSERT INTO songs(name, song, key, time) VALUES("{user}", "{request}", "{key}", "{time}");'''.format(user=user, request=(songtitle.replace('"', "'")), key=(key.replace('"', "'")), time=songtime))
                removetopqueue()
                return user + " >> Added: " + songtitle + " to the queue. ID: " + getnewentry()
        else:
            return self.youtubesr(request, user)

    def youtubesr(self, request, user):
        request = unidecode(request)
        try:
            results = (Mobileclient.search(api, request, SONGBLSIZE)['video_hits'])
            top_result = results[0]['youtube_video']
            video_url = "https://www.youtube.com/watch?v=" + top_result['id']
            key = top_result['id']
            title = top_result['title']
        except IndexError:
            try:
                self.video = YouTube("https://www.youtube.com/watch?v=" + request)
                return self.songrequest(" https://www.youtube.com/watch?v=" + request, user)
            except:
                return "No results at all"

        try:
            self.video = YouTube(video_url)
        except:
            return user + " >> That video is unavailable, it's probably age restricted."
        # Check the queue to see if the song is already in there.
        self.db = sqliteread('''SELECT id, count(*) FROM songs WHERE key="{0}"'''.format(key))
        if self.db[1] > (MAX_DUPLICATE_SONGS - 1):
            return user + " >> That song is already in the queue."
        songtime = self.getsongtime(video_url, key)
        if songtime > (MAXTIME * 60000):
            return user + " >> That song exceeds the maximum length of " + str(MAXTIME) + " minutes."

        sqlitewrite('''INSERT INTO songs(name, song, key, time) VALUES("{user}", "{request}", "{key}", "{time}");'''.format(user=user, request=(video_url.replace('"', "'")), key=(key.replace('"', "'")), time=songtime))
        removetopqueue()
        return user + " >> Added: " + title + " to the queue (YT). ID: " + getnewentry()

    '''--------------------QUEUE CONTROL--------------------'''

    def wrongsong(self, songid, user):
        if not songid:
            try:
                result = sqliteread('SELECT id, song FROM songs WHERE name="{0}"  ORDER BY id DESC LIMIT 1'.format(user))
                sqlitewrite('DELETE FROM songs WHERE id={0}'.format(str(result[0])))
                return user + ' >> Removed your request: "' + str(result[1]) + '" from the queue.'
            except Error as e:
                raise e
            except:
                return user + " >> Couldn't find your most recent request."
        else:
            try:
                result = sqliteread('SELECT song, name FROM songs WHERE id={0}'.format(songid))
                if user in result[1]:
                    sqlitewrite('DELETE FROM songs WHERE id={0}'.format(songid))
                    return user + ' >> Removed your request: "' + str(result[0]) + '" from the queue.'
                else:
                    return user + " >> You didn't request that song, you can't delete it!"
            except Error as e:
                raise e
            except:
                return user + " >> Couldn't find that request."

    def clearsong(self, songid, user):
        if not songid:
            try:
                result = sqliteread('SELECT id, song, name FROM songs ORDER BY id DESC LIMIT 1')
                sqlitewrite('DELETE FROM songs WHERE id={0}'.format(str(result[0])))
                return user + ' >> Removed the request: "' + str(result[1]) + '" requested by ' + str(result[2]) + " from the queue."
            except Error as e:
                raise e
            except:
                return "Something messed up. Is the queue empty?"
        else:
            try:
                result = sqliteread('SELECT song, name FROM songs WHERE id={0}'.format(songid))
                sqlitewrite('DELETE FROM songs WHERE id={0}'.format(songid))
                return user + ' >> Removed the request: "' + str(result[0]) + '" requested by ' + str(result[1]) +" from the queue."
            except Error as e:
                raise e
            except:
                return user + " >> Couldn't find that request."

    def queuetime(self, id, user):
        data = []
        db = sqlite3.connect('songqueue.db')
        try:
            cursor = db.cursor()
            if not id:  # If there's no ID, get the total song
                cursor.execute('''SELECT time FROM songs''')
            else:  # Get up to that song
                cursor.execute('''SELECT time FROM songs WHERE id < {0}'''.format(id))
            data = cursor.fetchall()
            if (not data) or (data[0][0] == None):
                if id:
                    return user + " >> That ID is not in the queue."
                else:
                    return user + " >> There are currently no songs in the queue."
            totaltime = 0
            for item in data:
                totaltime += int(item[0])
            seconds=(totaltime/1000)%60
            minutes=(totaltime/(1000*60))%60
            hours=(totaltime/(1000*60*60))%24
            db.close()
            return user + " >> There is about " + str(hours) + "h " + str(minutes) + "m " + str(seconds) + "s of songs in the queue."
        except Error as e:
            db.rollback()
            print ("QUEUETIME ERROR:")
            print (e)

    def skip(self, user, x):
        global skiprequests, skipusers
        from Run import veto

        if user in skipusers:
            return user + " >> You've already requested to skip this song."
        #Do the stuff to process the skip
        skipusers.append(user)
        skiprequests += 1

        if skiprequests >= VOTES_TO_SKIP:
            veto(None, None)
            return user + " >> Your vote to skip was enough to skip the song."

    '''--------------------BACKUP PL CONTROL--------------------'''

    def plsongrequest(self, request, user):
        request = request.replace("\r", "")[1:]
        # LINKS
        # DETERMINE LINK TYPE
        if validators.url(request):
            if "youtu" in request:
                try:
                    title = YouTube(request).title
                except Exception: return user + " >> Unable to use that video for some reason."
                key = getytkey(request)
                if not key: return user + " >> Something is wrong with your Youtube link."

                # Check the queue to see if the song is already in there.
                self.db = sqliteread('''SELECT id, count(*) FROM playlist WHERE key="{0}"'''.format(key))
                if self.db[1] > (MAX_DUPLICATE_SONGS - 1):
                    return user + " >> That song is already in the playlist. ID: " + str(self.db[0])
                sqlitewrite('''INSERT INTO playlist(song, key) VALUES("{request}", "{key}");'''.format(request=(request.replace('"', "'")), key=(key.replace('"', "'"))))
                return user + " >> Added: " + title + " to the playlist (YT). ID: " + str(sqliteread('SELECT id FROM playlist ORDER BY id DESC LIMIT 1')[0])
            else:  # OTHER MP3 REQUESTS <<<<<<<
                sqlitewrite('''INSERT INTO playlist(song, key) VALUES("{request}", "{request}");'''.format(request=(request.replace('"', "'"))))
                return user + " >> Added that link to the playlist. ID: " + str(sqliteread('SELECT id FROM playlist ORDER BY id DESC LIMIT 1')[0])
        else:  # GOOGLE PLAY MUSIC STUFF
            try:
                top_song_result = songtitlefilter(request, 0)
                key = top_song_result['storeId']
                songtitle = str(top_song_result['artist'] + " - " + top_song_result['title'])
            # If theres an error (its unable to find the song) then announce it, otherwise write the song data to the db
            except IndexError:
                return user + " >> No results found for that song. Please try a different one."
            else:
                # Test if the song is already in the playlist
                self.db = sqliteread('''SELECT id, count(*) FROM playlist WHERE key="{0}"'''.format(key))
                if self.db[1] > (MAX_DUPLICATE_SONGS - 1):
                    return user + " >> That song is already in the playlist. ID: " + str(self.db[0])
                # Add song to the playlist
                sqlitewrite('''INSERT INTO playlist(song, key) VALUES("{request}", "{key}");'''.format(request=(songtitle.replace('"', "'")), key=(key.replace('"', "'"))))
                return user + " >> Added: " + songtitle + " to the playlist. ID: " + str(sqliteread('SELECT id FROM playlist ORDER BY id DESC LIMIT 1')[0])

    def plclearsong(self, songid, user):
        if not songid:
            try:
                result = sqliteread('SELECT id, song FROM playlist ORDER BY id DESC LIMIT 1')
                sqlitewrite('DELETE FROM playlist WHERE id={0}'.format(str(result[0])))
                return user + ' >> Removed your request: "' + str(result[1]) + '" from the backup playlist.'
            except Exception as e:
                print (e)
                return user + " >> Couldn't find the most recent request."
        else:
            try:
                result = sqliteread('SELECT song FROM playlist WHERE id={0}'.format(songid))
                sqlitewrite('DELETE FROM playlist WHERE id={0}'.format(songid))
                return user + ' >> Removed the song: "' + str(result[0]) + " from the queue."
            except Error as e:
                raise e
            except:
                return user + " >> Couldn't find that request."

    '''--------------------MISC--------------------'''

    def getnowplaying(self, x, user):
        with open("Output/nowplaying.txt", "r") as f:
            returnnp = f.readlines()
        if not returnnp:
            sendMessage(s, (user + " >> The music is currently paused."))
        else:
            sendMessage(s, (user + " >> " + returnnp[0]))

    def clearqueue(self, x, y):
        try:
            sqlitewrite('DELETE FROM songs')
            return "Cleared the current songrequest queue"
        except:
            return "There was some sort of issue clearing the queue."

    def queuelink(self, user, x):
        return user + " >> " + QUEUE_LINK

    def getsongtime(self, title, key):
        songtime = -1
        try:
            if title:
                if title == "Online":
                    songurl = key
                else:
                    songurl = self.video.streams.filter(only_audio=True).order_by('abr').first().url
            else:  # Otherwise it's GPM
                songurl = sr_geturl(key)
            instance = vlc.Instance()
            media = instance.media_new(songurl)
            player = instance.media_player_new()
            player.set_media(media)
            #Start the parser
            media.parse_with_options(1,0)
            cycle = 0
            while True:
                cycle += 1
                if str(media.get_parsed_status()) == 'MediaParsedStatus.done':
                    break
                if cycle > 99999999:
                    print ("CRITICAL ERROR - GETSONGTIME IS BROKEN!")
                    break
            songtime = media.get_duration()
        # Let it be known to whatever brave adventurer is exploring my code, that this is the place that
        # Grant's sanity died for nearly two weeks straight.
        except Exception as e:
            print ("GETSONGTIME ISSUE")
            print (e)
            return
        return (songtime - 2000) #Most songs have a short offset.


'''
Schrijf een script dat gebruikt kan worden om getallen af te printen.
Het script ondersteunt de volgende argumenten:

    --init      Beginwaarde. Optioneel, standaardwaarde 0.
    --to        Eindwaarde. Verplicht.
    --step      Stapwaarde. Optioneel, standaardwaarde 1.

Voorbeeld 1

    python student.py --to 5

geeft

    0
    1
    2
    3
    4
    5

Voorbeeld 2

    python student.py --init 2 --to 10 --step 2

drukt af

    2
    4
    6
    8
    10

Steun op de module argparse (zoek informatie online en experimenteer)
om de command line arguments te parsen. Bijvoorbeeld,

    parser = argparse.ArgumentParser()
    parser.add_argument('--to', type=int, required=True)
    args = parser.parse_args()

zorgt ervoor dat je script een parameter "--to N" verwacht, waarbij
N een integer is. De variabele args bevat nadien een object
met veld 'to'. Bijvoorbeeld,

    python student.py --to 42

zorgt ervoor dat args.to gelijk is aan 42. Lees de documentatie
van argparse om te weten te komen hoe je optionele parameters kan instellen.
'''

import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--init', type=int)
    parser.add_argument('--to', type=int)
    parser.add_argument('--step', type=int)
    args = parser.parse_args()

    for i in range(args.init, args.to + 1, args.step):
        print(i)


from django import forms
from django.contrib.auth.forms import UserCreationForm
from django.contrib.auth.models import User
from django import forms
from django.contrib.auth.forms import AuthenticationForm
from django.utils.translation import gettext_lazy as _
from .models import Customer

class SignUpForm(UserCreationForm):
    date_of_birth = forms.DateField()

    class Meta:
        model = User
        fields = ('username', 'email', 'password1', 'password2', 'date_of_birth')

class RememberMeAuthenticationForm(AuthenticationForm):
    remember_me = forms.BooleanField(
        required=False,
        initial=True,
        widget=forms.CheckboxInput(attrs={'class': 'form-check-input'}),
        label=_('Remember me')
    )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if 'remember_me' in self.data:
            self.fields['remember_me'].initial = self.data['remember_me']

class CustomerForm(forms.ModelForm):
    class Meta:
        model = Customer
        fields = [
            'first_name',
            'middle_name',
            'last_name',
            'date_of_birth',
            'ssn',
            'citizenship',
            'country_of_residence',
            'address',
            'state',
            'zipcode',
            'passport',
            'utility_bill'
        ]


import torch
t = torch
import numpy as np
nn = t.nn
F = t.nn.functional
from transformers import BertJapaneseTokenizer, BertForMaskedLM
import datetime
from itertools import chain
from reader import read_data, read_regular_ds_zip
from sklearn.metrics import precision_recall_fscore_support as score
from sklearn.metrics import f1_score, precision_score, recall_score
import types

def create_model(learning_rate = 1e-5):
    res = types.SimpleNamespace()
    # NOTE
    res.bert = BertForMaskedLM.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
    res.toker = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
    res.opter = torch.optim.AdamW(res.bert.parameters(), learning_rate)
    res.bert.cuda()
    res.bert.train()
    return res

def verbalize(label):
    if label == 0:
        return 'いえ'
    elif label == 1:
        return 'はい'

def deverbalize(word):
    if word == 'いえ':
        return 0
    elif word == 'はい':
        return 1
    else:
        print(f'Bad word: {word}')
        return 0
        # return word

def patterned(left, right, strategy = 0):
    if strategy == 0:
        return f'「{left}」は「{right}」より上品ですか？[MASK]。'
    elif strategy == 1:
        return f'{left}は{right}よりも上品ですか？[MASK]。'
    elif strategy == 2:
        return f'「{left}」は「{right}」よりも格調高いですか？[MASK]。'
    elif strategy == 3:
        return f'{left}は{right}よりも格調高いですか？[MASK]。'

def pattern_verbalized(left, right, label):
    text = patterned(left, right)
    return text.replace('[MASK]', verbalize(label))

# NOTE: create half fake datas
def create_inputs_and_labels_from_fewshot_set(ds, should_deverbalize = True, shuffle = True):
    results = []
    ds = ds.copy()
    if shuffle:
        np.random.shuffle(ds)
    half = int(len(ds) / 2)
    ds_true = ds[:half]
    ds_false = ds[half:]
    for left, right in ds_true:
        results.append((patterned(left, right), pattern_verbalized(left, right, 0) if should_deverbalize else 0))
    for left, right in ds_false:
        results.append((patterned(right, left), pattern_verbalized(right, left, 1) if should_deverbalize else 1))
    if shuffle:
        np.random.shuffle(results)
    return results

def step(x, y, m):
    tokenizer = m.toker
    model = m
    opter = m.opter
    inputs = tokenizer(x, return_tensors="pt", truncation=True)["input_ids"]
    labels = tokenizer(y, return_tensors="pt", truncation=True)["input_ids"]
    labels = torch.where(inputs == tokenizer.mask_token_id, labels, -100)
    inputs = inputs.cuda()
    labels = labels.cuda()
    outputs = model.bert(inputs, labels=labels)
    loss = outputs.loss
    loss.backward()

def run_full(m, ds, batch_size = 1):
    input_labels = create_inputs_and_labels_from_fewshot_set(ds)
    length = len(input_labels)
    first_time = datetime.datetime.now()
    for index, (x,y) in enumerate(input_labels):
        step(x, y, m)
        if (index + 1) % 32 == 0:
            print(f'{index + 1} / {length}')
        if (index + 1) % batch_size == 0:
            m.opter.step()
            m.opter.zero_grad()
    m.opter.step()
    m.opter.zero_grad()
    last_time = datetime.datetime.now()
    delta = last_time - first_time
    print(delta.seconds)

def get_predicted_word(tokenizer, model, text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True)["input_ids"]
    inputs = inputs.cuda()
    with torch.no_grad():
        logits = model(inputs).logits
    mask_token_index = (inputs == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
    predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
    word = tokenizer.decode(predicted_token_id)
    return word

def get_test_result(m, ds_test, should_deverbalize = True):
    tokenizer = m.toker
    model = m.bert
    pred_ys = []
    true_ys = []
    length = len(ds_test)
    input_labels = create_inputs_and_labels_from_fewshot_set(ds_test, should_deverbalize = False, shuffle = False)
    # for index, item in enumerate(test):
    for index, (text, label) in enumerate(input_labels):
        if index % 100 == 0:
            print(f'{index}/{length}')
        word = get_predicted_word(tokenizer, model, text)
        pred_y = deverbalize(word) if should_deverbalize else word
        pred_ys.append(pred_y)
        true_ys.append(label)
    return pred_ys, true_ys

def cal_scores(pred_ys, true_ys):
    precision, recall, fscore, support = score(true_ys, pred_ys)

def cal_prec_rec_f1_v2(results, targets):
    return f1_score(targets, results, average='macro')

def get_curve_regular():
    m = create_model()
    ds = read_regular_ds_zip()
    train_ds = ds[:227] # 227
    test_ds = ds[227:277] # 50
    fs = []
    precs = []
    recs = []
    for start in range(0, 224, 16): # [0, 16, 32, ..., 480]
        end = start + 16 # 16, 32, ..., 496
        fewshot = train_ds[start: end]
        run_full(m, fewshot, batch_size = 1)
        results, targets = get_test_result(m, test_ds)
        fs.append(f1_score(targets, results, average='macro'))
        precs.append(precision_score(targets, results, average='macro'))
        recs.append(recall_score(targets, results, average='macro'))
    return precs, recs, fs

def get_curve_regular(m, batch_size = 1):
    ds = read_regular_ds_zip()
    train_ds = ds[:227] # 227
    test_ds = ds[227:277] # 50
    # ds = read_data()
    # train_ds = ds[:496]
    # test_ds = ds[500:700]
    fs = []
    precs = []
    recs = []
    # for start in range(0, 496, 32): # [0, 16, 32, ..., 480]
    for start in range(0, 227, 32): # [0, 16, 32, ..., 480]
        end = start + 32 # 16, 32, ..., 496
        fewshot = train_ds[start: end]
        run_full(m, fewshot, batch_size = batch_size)
        results, targets = get_test_result(m, test_ds)
        fs.append(f1_score(targets, results, average='macro'))
        precs.append(precision_score(targets, results, average='macro'))
        recs.append(recall_score(targets, results, average='macro'))
    return precs, recs, fs

def get_curve(m, batch_size = 1):
    ds = read_data()
    train_ds = ds[:496]
    test_ds = ds[500:700]
    fs = []
    precs = []
    recs = []
    for start in range(0, 496, 32): # [0, 16, 32, ..., 480]
        end = start + 32 # 16, 32, ..., 496
        fewshot = train_ds[start: end]
        run_full(m, fewshot, batch_size = batch_size)
        results, targets = get_test_result(m, test_ds)
        fs.append(f1_score(targets, results, average='macro'))
        precs.append(precision_score(targets, results, average='macro'))
        recs.append(recall_score(targets, results, average='macro'))
    return precs, recs, fs

def stop_when_f_score_exceed(limit = 0.57):
    m = create_model()
    ds = read_data()
    train_ds = ds[:496]
    test_ds = ds[500:700]
    fs = []
    precs = []
    recs = []
    for start in range(0, 496, 16): # [0, 16, 32, ..., 480]
        end = start + 16 # 16, 32, ..., 496
        fewshot = train_ds[start: end]
        run_full(m, fewshot, batch_size = 1)
        results, targets = get_test_result(m, test_ds)
        f = f1_score(targets, results, average='macro')
        if f > limit:
            print(f'End at {start}, f {f}.')
            break
    return m

# ========================================

def train_model(data_points = 256, batch_size = 1):
    m = create_model()
    ds = read_data()
    assert data_points < 500
    train_ds = ds[:data_points]
    test_ds = ds[500:700]
    print('Training...')
    run_full(m, train_ds, batch_size = batch_size)
    print('Testing...')
    results, targets = get_test_result(m, test_ds)
    f = f1_score(targets, results, average='macro')
    prec = precision_score(targets, results, average='macro')
    rec = recall_score(targets, results, average='macro')
    print(f'Results: f {f}, precision {prec}, recall {rec}.')
    return m

def get_predicted_word_by_word_pair(m, left, right):
    text = patterned(left, right)
    return get_predicted_word(m.toker, m.bert, text)

def run_test(m, test_ds):
    results, targets = get_test_result(m, test_ds)
    print(f1_score(targets, results, average='macro'))


import torch
from torch.utils.data import Dataset
import json
import numpy as np
import os
from PIL import Image
from torchvision import transforms as T
import time

from ray_utils import *
from einops import rearrange,repeat
import torch.nn.functional as F

class BlenderDataset(Dataset):
    def __init__(self, root_dir, split='train', img_wh=(800, 800)):
        self.root_dir = root_dir
        self.split = split
        assert img_wh[0] == img_wh[1], 'image width must equal image height!'
        self.img_wh = img_wh
        self.transforms = T.ToTensor()

        self.read_meta()
        self.white_back = True

    def read_meta(self):
        with open(os.path.join(self.root_dir,
                               f"transforms_{self.split}.json"), 'r') as f:
            self.meta = json.load(f)

        self.w, self.h = self.img_wh
        self.focal = 0.5*800/np.tan(0.5*self.meta['camera_angle_x']) # original focal length
                                                                     # when W=800

        self.focal *= self.img_wh[0]/800 # modify focal length to match size self.img_wh

        # bounds, common for all scenes
        self.near = 2.0
        self.far = 6.0

        self.directions = get_ray_directions(self.h, self.w, self.focal) # (h, w, 3)

        if self.split == 'train':
            self.image_paths = []
            self.poses = []
            self.all_ray_o = []
            self.all_ray_d = []
            self.all_rgbs = []
            self.viewdirs = []
            print('get_training_rays: start')
            eps_time = time.time()
            for frame in self.meta['frames']:
                pose = np.array(frame['transform_matrix'])[:3, :4]
                self.poses += [pose]
                c2w = torch.FloatTensor(pose)

                image_path = os.path.join(self.root_dir, f"{frame['file_path']}.png")
                self.image_paths += [image_path]
                img = Image.open(image_path)
                img = img.resize(self.img_wh, Image.LANCZOS)
                img = self.transforms(img) # (4, h, w)
                img = img.view(4, -1).permute(1, 0) # (h*w, 4) RGBA
                img = img[:, :3]*img[:, -1:] + (1-img[:, -1:]) # blend A to RGB
                self.all_rgbs += [img]

                rays_o, rays_d = get_rays(self.directions, c2w)  # both (h*w, 3)

                #rays_o, rays_d = get_view_rays(800, 800, pose_c2w = c2w, focal = self.focal, pixel_mode='center')
                #rays_o = rays_o.reshape(-1,3)
                #rays_d = rays_d.reshape(-1,3)

                viewdirs = rays_d / rays_d.norm(dim=-1, keepdim=True)

                self.all_ray_o += [rays_o]
                self.all_ray_d += [rays_d]
                self.viewdirs += [viewdirs]

            self.all_ray_o = torch.cat(self.all_ray_o, 0) # (len(self.meta['frames])*h*w, 3)
            self.all_ray_d = torch.cat(self.all_ray_d, 0) # (len(self.meta['frames])*h*w, 3)
            self.all_rgbs = torch.cat(self.all_rgbs, 0) # (len(self.meta['frames])*h*w, 3)
            self.viewdirs = torch.cat(self.viewdirs, 0)  # (len(self.meta['frames])*h*w, 3)
            print('get_training_rays: finish (eps time:', eps_time, 'sec)')

    def get_bound(self):
        return self.near, self.far

    def define_bbox(self):
        print('Compute_bbox: start')
        xyz_min = torch.Tensor([np.inf, np.inf, np.inf])
        xyz_max = -xyz_min
        pts_nf = torch.stack([self.all_ray_o+self.viewdirs*self.near, self.all_ray_o+self.viewdirs*self.far])
        xyz_min = torch.minimum(xyz_min, pts_nf.amin((0,1)))
        xyz_max = torch.maximum(xyz_max, pts_nf.amax((0,1)))
        print('Compute_bbox: xyz_min', xyz_min)
        print('Compute_bbox: xyz_max', xyz_max)
        print('Compute_bbox: finish')
        return xyz_max, xyz_min

    def voxel_count_views(self, channel, Nx, Ny, Nz, xyz_max, xyz_min, voxel_size):
        print('dvgo: voxel_count_views start')
        eps_time = time.time()
        count = torch.zeros([1, channel, Nx, Ny, Nz]).to('cuda')
        samples_per_ray = int(np.linalg.norm(np.array([Nx, Ny, Nz])+1) / 0.5) + 1
        z_vals = torch.arange(samples_per_ray).to('cuda')*voxel_size*0.5
        chunk_size = 10000

        rays_o = self.all_ray_o.to('cuda').split(chunk_size)
        rays_d = self.all_ray_d.to('cuda').split(chunk_size)

        ones = torch.ones([1, channel, Nx, Ny, Nz]).to('cuda').requires_grad_()
        for ray_os, ray_ds in zip(rays_o, rays_d):
            ray_d_norm = ray_ds.norm(dim=-1, keepdim=True)
            #viewdirs = ray_ds/ray_d_norm

            t_1 = ((xyz_max.to(ray_os.device) - ray_os) / (ray_ds + 1e-6)) #(num_rays,3)
            t_2 = ((xyz_min.to(ray_os.device) - ray_os) / (ray_ds + 1e-6)) #(num_rays,3)

            t_min = torch.minimum(t_1, t_2).amax(-1).clamp(min=self.near, max=self.far)
            t_max = torch.maximum(t_1, t_2).amin(-1).clamp(min=self.near, max=self.far)

            delta = torch.einsum('ij,ik->ikj', [ray_ds/ray_d_norm, t_min[...,None]*ray_d_norm  + z_vals[None,:]])
            rays_pts = ray_os[...,None, :] + delta
            '''
            ray_start = ray_os + t_min[...,None] * ray_ds
            rays_pts = ray_start[...,
                      None, :] + z_vals[..., :, None] * viewdirs[..., None, :]
            '''
            xyz_min_ = xyz_min.to(rays_pts.device).reshape(1, 1, -1)
            xyz_max_ = xyz_max.to(rays_pts.device).reshape(1, 1, -1)
            rays_pts = ((rays_pts - xyz_min_)/(xyz_max_ - xyz_min_))*2 - 1
            rays_pts = rays_pts.reshape(1, 1, 1, -1, 3)
            output = F.grid_sample(ones, rays_pts, mode='bilinear', align_corners=True)
            output = output.sum().backward()

            with torch.no_grad():
                count += (ones.grad > 1)

        eps_time = time.time() - eps_time
        print('dvgo: voxel_count_views finish (eps time:', eps_time, 'sec)')
        return count

    def __len__(self):
        if self.split == 'train':
            return len(self.all_rgbs)
        if self.split == 'val':
            return 8 # only validate 8 images (to support <=8 gpus)
        return len(self.meta['frames'])

    def __getitem__(self, idx):
        if self.split == 'train':
            sample = {  'ray_o': self.all_ray_o[idx],
                    'ray_d': self.all_ray_d[idx],
                    'viewdirs' : self.viewdirs[idx],
                    'rgbs': self.all_rgbs[idx]}

        else: # create data for each image separately
            frame = self.meta['frames'][idx]
            c2w = torch.FloatTensor(frame['transform_matrix'])[:3, :4]

            img = Image.open(os.path.join(self.root_dir, f"{frame['file_path']}.png"))
            img = img.resize(self.img_wh, Image.LANCZOS)
            img = self.transforms(img) # (4, H, W)
            valid_mask = (img[-1]>0).flatten() # (H*W) valid color area
            img = img.view(4, -1).permute(1, 0) # (H*W, 4) RGBA
            img = img[:, :3]*img[:, -1:] + (1-img[:, -1:]) # blend A to RGB

            rays_o, rays_d = get_rays(self.directions, c2w)

            '''
            rays_o, rays_d = get_view_rays(800, 800, pose_c2w = c2w, focal = self.focal, pixel_mode='center')
            rays_o = rays_o.reshape(-1,3)
            rays_d = rays_d.reshape(-1,3)
            '''

            viewdirs = rays_d / rays_d.norm(dim=-1, keepdim=True)

            sample = {'ray_o': rays_o,
                      'ray_d': rays_d,
                      'viewdirs' : viewdirs,
                      'rgbs': img,
                      'c2w': c2w,
                      'valid_mask': valid_mask}

        return sample

if __name__ == "__main__":
    dataset = BlenderDataset(
            root_dir = 'data/nerf_synthetic/lego', split='val', img_wh=(800, 800)
        )

    from torch.utils.data import DataLoader
    dataloader = DataLoader(
            dataset,
            shuffle=False,
            num_workers=1,
            batch_size = 1,
            pin_memory=True
        )

    for data in dataloader:
        rays_o = data['rays_o']
        rays_d = data['rays_d']
        #print(rays_o.shape)
        #print(rays_d.shape)
        break


from playwright.sync_api import sync_playwright

# 'url' is the URL that playwright's browser shpuld visit
# 'session_data' is the data stored in your own session, this needs to be transferred to be worked with
def take_screenshot_from_url(url, session_data):
    with sync_playwright() as playwright:   # Using a context manager to wrap playwright code, ensures python sets up and tears down playwright instance properly
        webkit = playwright.webkit  # Create headless webkit browser
        browser = webkit.launch()   # launch the webkit browser
        browser_context = browser.new_context(device_scale_factor=2)    # set up new context for browser, and increae deviec_scale_factor at 2 to ensure pictures aren't pixelated
        browser_context.add_cookies([session_data]) # Add session info to the browser
        page = browser_context.new_page()   # open a new page in the browser_context, like a new browser window
        page.goto(url)  # visit the target url
        screenshot_bytes = page.locator(".code").screenshot()   # Take a ss of the image, this returns a buffer with the image
        browser.close()
        return screenshot_bytes # Return the image buffer