"""
This module contains the implementation of a verification setup for a poromechanical
system (without fractures) using two-dimensional manufactured solutions.

The implementation considers a compressible fluid, resulting in a coupled non-linear
set of equations. The dependecy of the fluid density with the pressure is given by [1]:

.. math::

    \\rho(p) =  \\rho_0 \\exp{c_f \\left(p - p_0\\right)},

where :math:`\\rho` and :math:`p` are the density and pressure, :math:`\\rho_0`` and
:math:``p_0`` are the _reference_ density and pressure, and :math:`c_f` is the
_constant_ fluid compressibility.

We provide two different manufactured solutions, namely `parabolic` and `nordbotten2016`
[2], which can be specified as a model parameter.

Parabolic manufactured solution:

.. math::

    p(x, y, t) = t x (1 - x) y (1- y),

    u_x(x, y, t) = p(x, y, t),

    u_y(x, y, t) = p(x, y, t).

Manufactured solution based on [2]:

.. math::

    p(x, y, t) = t x (1 - x) \\sin{2 \\pi y},

    u_x(x, y, t) = p(x, y, t),

    u_y(x, y, t) = t * \\sin{2 \\pi x} \\sin{2 \\pi y}.

References:

    - [1] Garipov, T. T., & Hui, M. H. (2019). Discrete fracture modeling approach for
      simulating coupled thermo-hydro-mechanical effects in fractured reservoirs.
      International Journal of Rock Mechanics and Mining Sciences, 122, 104075.

    - [2] Nordbotten, J. M. (2016). Stable cell-centered finite volume discretization
      for Biot equations. SIAM Journal on Numerical Analysis, 54(2), 942-968.

"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np
import sympy as sym

import porepy as pp
import porepy.models.fluid_mass_balance as mass
import porepy.models.momentum_balance as momentum
import porepy.models.poromechanics as poromechanics
from porepy.applications.convergence_analysis import ConvergenceAnalysis
from porepy.applications.md_grids.domains import nd_cube_domain
from porepy.utils.examples_utils import VerificationUtils
from porepy.viz.data_saving_model_mixin import VerificationDataSaving

# PorePy typings
number = pp.number
grid = pp.GridLike

# -----> Data-saving
@dataclass
class ManuPoroMechSaveData:
    """Data class to save relevant results from the verification setup."""

    approx_displacement: np.ndarray
    """Numerical displacement."""

    approx_flux: np.ndarray
    """Numerical Darcy flux."""

    approx_force: np.ndarray
    """Numerical poroelastic force."""

    approx_pressure: np.ndarray
    """Numerical pressure."""

    error_displacement: number
    """L2-discrete relative error for the displacement."""

    error_flux: number
    """L2-discrete relative error for the Darcy flux."""

    error_force: number
    """L2-discrete relative error for the poroelastic force."""

    error_pressure: number
    """L2-discrete relative error for the pressure."""

    exact_displacement: np.ndarray
    """Exact displacement."""

    exact_flux: np.ndarray
    """Exact Darcy flux."""

    exact_force: np.ndarray
    """Exact poroelastic force."""

    exact_pressure: np.ndarray
    """Exact pressure."""

    time: number
    """Current simulation time."""

class ManuPoroMechDataSaving(VerificationDataSaving):
    """Mixin class to save relevant data."""

    darcy_flux: Callable[[list[pp.Grid]], pp.ad.Operator]
    """Method that returns the Darcy fluxes in the form of an Ad operator. Usually
    provided by the mixin class :class:`porepy.models.constitutive_laws.DarcysLaw`.

    """

    displacement: Callable[[list[pp.Grid]], pp.ad.MixedDimensionalVariable]
    """Displacement variable. Normally defined in a mixin instance of
    :class:`~porepy.models.momentum_balance.VariablesMomentumBalance`.

    """

    exact_sol: ManuPoroMechExactSolution2d
    """Exact solution object."""

    pressure: Callable[[list[pp.Grid]], pp.ad.MixedDimensionalVariable]
    """Pressure variable. Normally defined in a mixin instance of
    :class:`~porepy.models.fluid_mass_balance.VariablesSinglePhaseFlow`.

    """

    stress: Callable[[list[pp.Grid]], pp.ad.Operator]
    """Method that returns the (integrated) poroelastic stress in the form of an Ad
    operator. Usually provided by the mixin class
    :class:`porepy.models.poromechanics.ConstitutiveLawsPoromechanics`.

    """

    def collect_data(self) -> ManuPoroMechSaveData:
        """Collect data from the verification setup.

        Returns:
            ManuPoroMechSaveData object containing the results of the verification for
            the current time.

        """

        mdg: pp.MixedDimensionalGrid = self.mdg
        sd: pp.Grid = mdg.subdomains()[0]
        t: number = self.time_manager.time

        # Collect data
        exact_pressure = self.exact_sol.pressure(sd=sd, time=t)
        pressure_ad = self.pressure([sd])
        approx_pressure = pressure_ad.value(self.equation_system)
        error_pressure = ConvergenceAnalysis.l2_error(
            grid=sd,
            true_array=exact_pressure,
            approx_array=approx_pressure,
            is_scalar=True,
            is_cc=True,
            relative=True,
        )

        exact_displacement = self.exact_sol.displacement(sd=sd, time=t)
        displacement_ad = self.displacement([sd])
        approx_displacement = displacement_ad.value(self.equation_system)
        error_displacement = ConvergenceAnalysis.l2_error(
            grid=sd,
            true_array=exact_displacement,
            approx_array=approx_displacement,
            is_scalar=False,
            is_cc=True,
            relative=True,
        )

        exact_flux = self.exact_sol.darcy_flux(sd=sd, time=t)
        flux_ad = self.darcy_flux([sd])
        approx_flux = flux_ad.value(self.equation_system)
        error_flux = ConvergenceAnalysis.l2_error(
            grid=sd,
            true_array=exact_flux,
            approx_array=approx_flux,
            is_scalar=True,
            is_cc=False,
            relative=True,
        )

        exact_force = self.exact_sol.poroelastic_force(sd=sd, time=t)
        force_ad = self.stress([sd])
        approx_force = force_ad.value(self.equation_system)
        error_force = ConvergenceAnalysis.l2_error(
            grid=sd,
            true_array=exact_force,
            approx_array=approx_force,
            is_scalar=False,
            is_cc=False,
            relative=True,
        )

        # Store collected data in data class
        collected_data = ManuPoroMechSaveData(
            approx_displacement=approx_displacement,
            approx_flux=approx_flux,
            approx_force=approx_force,
            approx_pressure=approx_pressure,
            error_displacement=error_displacement,
            error_flux=error_flux,
            error_force=error_force,
            error_pressure=error_pressure,
            exact_displacement=exact_displacement,
            exact_flux=exact_flux,
            exact_force=exact_force,
            exact_pressure=exact_pressure,
            time=t,
        )

        return collected_data

# -----> Exact solution
class ManuPoroMechExactSolution2d:
    """Class containing the exact manufactured solution for the verification setup."""

    def __init__(self, setup):
        """Constructor of the class."""

        # Physical parameters
        lame_lmbda = setup.solid.lame_lambda()  # [Pa] Lamé parameter
        lame_mu = setup.solid.shear_modulus()  # [Pa] Lamé parameter
        alpha = setup.solid.biot_coefficient()  # [-] Biot coefficient
        rho_0 = setup.fluid.density()  # [kg / m^3] Reference density
        phi_0 = setup.solid.porosity()  # [-] Reference porosity
        p_0 = setup.fluid.pressure()  # [Pa] Reference pressure
        c_f = setup.fluid.compressibility()  # [Pa^-1] Fluid compressibility
        k = setup.solid.permeability()  # [m^2] Permeability
        mu_f = setup.fluid.viscosity()  # [Pa * s] Fluid viscosity
        K_d = lame_lmbda + (2 / 3) * lame_mu  # [Pa] Bulk modulus

        # Symbolic variables
        x, y, t = sym.symbols("x y t")
        pi = sym.pi

        # Exact pressure and displacement solutions
        manufactured_sol = setup.params.get("manufactured_solution", "parabolic")
        if manufactured_sol == "parabolic":
            p = t * x * (1 - x) * y * (1 - y)
            u = [p, p]
        elif manufactured_sol == "nordbotten_2016":
            p = t * x * (1 - x) * sym.sin(2 * pi * y)
            u = [p, t * sym.sin(2 * pi * x) * sym.sin(2 * pi * y)]
        else:
            raise NotImplementedError("Manufactured solution is not available.")

        # Exact density
        rho = rho_0 * sym.exp(c_f * (p - p_0))

        # Exact darcy flux
        q = [-1 * (k / mu_f) * sym.diff(p, x), -1 * (k / mu_f) * sym.diff(p, y)]

        # Exact mass flux
        mf = [rho * q[0], rho * q[1]]

        # Exact divergence of the mass flux
        div_mf = sym.diff(mf[0], x) + sym.diff(mf[1], y)

        # Exact divergence of the displacement
        div_u = sym.diff(u[0], x) + sym.diff(u[1], y)

        # Exact poromechanical porosity
        phi = phi_0 + ((alpha - phi_0) * (1 - alpha) / K_d) * (p - p_0) + alpha * div_u

        # Exact flow accumulation
        accum_flow = sym.diff(phi * rho, t)

        # Exact flow source
        source_flow = accum_flow + div_mf

        # Exact gradient of the displacement
        grad_u = [
            [sym.diff(u[0], x), sym.diff(u[0], y)],
            [sym.diff(u[1], x), sym.diff(u[1], y)],
        ]

        # Exact transpose of the gradient of the displacement
        trans_grad_u = [[grad_u[0][0], grad_u[1][0]], [grad_u[0][1], grad_u[1][1]]]

        # Exact strain
        epsilon = [
            [
                0.5 * (grad_u[0][0] + trans_grad_u[0][0]),
                0.5 * (grad_u[0][1] + trans_grad_u[0][1]),
            ],
            [
                0.5 * (grad_u[1][0] + trans_grad_u[1][0]),
                0.5 * (grad_u[1][1] + trans_grad_u[1][1]),
            ],
        ]

        # Exact trace (in the linear algebra sense) of the strain
        tr_epsilon = epsilon[0][0] + epsilon[1][1]

        # Exact elastic stress
        sigma_elas = [
            [
                lame_lmbda * tr_epsilon + 2 * lame_mu * epsilon[0][0],
                2 * lame_mu * epsilon[0][1],
            ],
            [
                2 * lame_mu * epsilon[1][0],
                lame_lmbda * tr_epsilon + 2 * lame_mu * epsilon[1][1],
            ],
        ]

        # Exact poroelastic stress
        sigma_total = [
            [sigma_elas[0][0] - alpha * p, sigma_elas[0][1]],
            [sigma_elas[1][0], sigma_elas[1][1] - alpha * p],
        ]

        # Mechanics source term
        source_mech = [
            sym.diff(sigma_total[0][0], x) + sym.diff(sigma_total[0][1], y),
            sym.diff(sigma_total[1][0], x) + sym.diff(sigma_total[1][1], y),
        ]

        # Public attributes
        self.p = p  # pressure
        self.u = u  # displacement
        self.sigma_elast = sigma_elas  # elastic stress
        self.sigma_total = sigma_total  # poroelastic (total) stress
        self.q = q  # Darcy flux
        self.mf = mf  # Mass flux, i.e., density * Darcy flux
        self.phi = phi  # poromechanics porosity
        self.rho = rho  # density
        self.source_mech = source_mech  # Source term entering the momentum balance
        self.source_flow = source_flow  # Source term entering the mass balance

    # -----> Primary and secondary variables
    def pressure(self, sd: pp.Grid, time: float) -> np.ndarray:
        """Evaluate exact pressure [Pa] at the cell centers.

        Parameters:
            sd: Subdomain grid.
            time: Time in seconds.

        Returns:
            Array of ``shape=(sd.num_cells, )`` containing the exact pressures at the
            cell centers at the given ``time``.

        """
        # Symbolic variables
        x, y, t = sym.symbols("x y t")

        # Get cell centers
        cc = sd.cell_centers

        # Lambdify expression
        p_fun: Callable = sym.lambdify((x, y, t), self.p, "numpy")

        # Cell-centered pressures
        p_cc: np.ndarray = p_fun(cc[0], cc[1], time)

        return p_cc

    def displacement(self, sd: pp.Grid, time: float) -> np.ndarray:
        """Evaluate exact displacement [m] at the cell centers.

        Parameters:
            sd: Subdomain grid.
            time: Time in seconds.

        Returns:
            Array of ``shape=(2 * sd.num_cells, )`` containing the exact displacements
            at the cell centers for the given ``time``.

        Notes:
            The returned displacement is given in PorePy's flattened vector format.

        """
        # Symbolic variables
        x, y, t = sym.symbols("x y t")

        # Get cell centers
        cc = sd.cell_centers

        # Lambdify expression
        u_fun: list[Callable] = [
            sym.lambdify((x, y, t), self.u[0], "numpy"),
            sym.lambdify((x, y, t), self.u[1], "numpy"),
        ]

        # Cell-centered displacements
        u_cc: list[np.ndarray] = [
            u_fun[0](cc[0], cc[1], time),
            u_fun[1](cc[0], cc[1], time),
        ]

        # Flatten array
        u_flat: np.ndarray = np.asarray(u_cc).ravel("F")

        return u_flat

    def darcy_flux(self, sd: pp.Grid, time: float) -> np.ndarray:
        """Evaluate exact Darcy flux [m^3 * s^-1] at the face centers.

        Parameters:
            sd: Subdomain grid.
            time: Time in seconds.

        Returns:
            Array of ``shape=(sd.num_faces, )`` containing the exact Darcy fluxes at
            the face centers for the given ``time``.

        Note:
            The returned fluxes are already scaled with ``sd.face_normals``.

        """
        # Symbolic variables
        x, y, t = sym.symbols("x y t")

        # Get list of face indices
        fc = sd.face_centers
        fn = sd.face_normals

        # Lambdify expression

        q_fun: list[Callable] = [
            sym.lambdify((x, y, t), self.q[0], "numpy"),
            sym.lambdify((x, y, t), self.q[1], "numpy"),
        ]

        # Face-centered Darcy fluxes
        q_fc: np.ndarray = (
            q_fun[0](fc[0], fc[1], time) * fn[0] + q_fun[1](fc[0], fc[1], time) * fn[1]
        )

        return q_fc

    def poroelastic_force(self, sd: pp.Grid, time: float) -> np.ndarray:
        """Evaluate exact poroelastic force [N] at the face centers.

        Parameters:
            sd: Subdomain grid.
            time: Time in seconds.

        Returns:
            Array of ``shape=(2 * sd.num_faces, )`` containing the exact poroealstic
            force at the face centers for the given ``time``.

        Notes:
            - The returned poroelastic force is given in PorePy's flattened vector
              format.
            - Recall that force = (stress dot_prod unit_normal) * face_area.

        """
        # Symbolic variables
        x, y, t = sym.symbols("x y t")

        # Get cell centers and face normals
        fc = sd.face_centers
        fn = sd.face_normals

        # Lambdify expression
        sigma_total_fun: list[list[Callable]] = [
            [
                sym.lambdify((x, y, t), self.sigma_total[0][0], "numpy"),
                sym.lambdify((x, y, t), self.sigma_total[0][1], "numpy"),
            ],
            [
                sym.lambdify((x, y, t), self.sigma_total[1][0], "numpy"),
                sym.lambdify((x, y, t), self.sigma_total[1][1], "numpy"),
            ],
        ]

        # Face-centered poroelastic force
        force_total_fc: list[np.ndarray] = [
            # (sigma_xx * n_x + sigma_xy * n_y) * face_area
            sigma_total_fun[0][0](fc[0], fc[1], time) * fn[0]
            + sigma_total_fun[0][1](fc[0], fc[1], time) * fn[1],
            # (sigma_yx * n_x + sigma_yy * n_y) * face_area
            sigma_total_fun[1][0](fc[0], fc[1], time) * fn[0]
            + sigma_total_fun[1][1](fc[0], fc[1], time) * fn[1],
        ]

        # Flatten array
        force_total_flat: np.ndarray = np.asarray(force_total_fc).ravel("F")

        return force_total_flat

    # -----> Sources
    def mechanics_source(self, sd: pp.Grid, time: float) -> np.ndarray:
        """Compute exact source term for the momentum balance equation.

        Parameters:
            sd: Subdomain grid.
            time: Time in seconds.

        Returns:
            Exact right hand side of the momentum balance equation with ``shape=(
            2 * sd.num_cells, )``.

        Notes:
            The returned array is given in PorePy's flattened vector format.

        """
        # Symbolic variables
        x, y, t = sym.symbols("x y t")

        # Get cell centers and cell volumes
        cc = sd.cell_centers
        vol = sd.cell_volumes

        # Lambdify expression
        source_mech_fun: list[Callable] = [
            sym.lambdify((x, y, t), self.source_mech[0], "numpy"),
            sym.lambdify((x, y, t), self.source_mech[1], "numpy"),
        ]

        # Evaluate and integrate source
        source_mech: list[np.ndarray] = [
            source_mech_fun[0](cc[0], cc[1], time) * vol,
            source_mech_fun[1](cc[0], cc[1], time) * vol,
        ]

        # Flatten array
        source_mech_flat: np.ndarray = np.asarray(source_mech).ravel("F")

        # Invert sign according to sign convention.
        return -source_mech_flat

    def flow_source(self, sd: pp.Grid, time: float) -> np.ndarray:
        """Compute exact source term for the fluid mass balance equation.

        Parameters:
            sd: Subdomain grid.
            time: Time in seconds.

        Returns:
            Exact right hand side of the fluid mass balance equation with ``shape=(
            sd.num_cells, )``.

        """
        # Symbolic variables
        x, y, t = sym.symbols("x y t")

        # Get cell centers and cell volumes
        cc = sd.cell_centers
        vol = sd.cell_volumes

        # Lambdify expression
        source_flow_fun: Callable = sym.lambdify((x, y, t), self.source_flow, "numpy")

        # Evaluate and integrate source
        source_flow: np.ndarray = source_flow_fun(cc[0], cc[1], time) * vol

        return source_flow

# -----> Utilities
class ManuPoroMechUtils(VerificationUtils):
    """Mixin class containing useful utility methods for the setup."""

    mdg: pp.MixedDimensionalGrid
    """Mixed-dimensional grid."""

    results: list[ManuPoroMechSaveData]
    """List of ManuPoroMechSaveData objects."""

    def plot_results(self) -> None:
        """Plotting results."""
        self._plot_displacement()
        self._plot_pressure()

    def _plot_pressure(self):
        """Plot exact and numerical pressures."""

        sd = self.mdg.subdomains()[0]
        p_ex = self.results[-1].exact_pressure
        p_num = self.results[-1].approx_pressure

        pp.plot_grid(sd, p_ex, plot_2d=True, linewidth=0, title="p (Exact)")
        pp.plot_grid(sd, p_num, plot_2d=True, linewidth=0, title="p (Numerical)")

    def _plot_displacement(self):
        """Plot exact and numerical displacements."""

        sd = self.mdg.subdomains()[0]
        u_ex = self.results[-1].exact_displacement
        u_num = self.results[-1].approx_displacement

        # Horizontal displacement
        pp.plot_grid(sd, u_ex[::2], plot_2d=True, linewidth=0, title="u_x (Exact)")
        pp.plot_grid(sd, u_num[::2], plot_2d=True, linewidth=0, title="u_x (Numerical)")

        # Vertical displacement
        pp.plot_grid(sd, u_ex[1::2], plot_2d=True, linewidth=0, title="u_y (Exact)")
        pp.plot_grid(
            sd, u_num[1::2], plot_2d=True, linewidth=0, title="u_y (Numerical)"
        )

# -----> Geometry
class UnitSquareGrid:
    """Class for setting up the geometry of the unit square domain."""

    params: dict
    """Simulation model parameters."""

    def set_domain(self) -> None:
        """Set domain."""
        self._domain = nd_cube_domain(2, 1.0)

    def meshing_arguments(self) -> dict[str, float]:
        """Set meshing arguments."""
        default_mesh_arguments = {"cell_size": 0.1}
        return self.params.get("meshing_arguments", default_mesh_arguments)

# -----> Balance equations
class ManuPoroMechMassBalance(mass.MassBalanceEquations):
    """Modify balance equation to account for external sources."""

    def fluid_source(self, subdomains: list[pp.Grid]) -> pp.ad.Operator:
        """Fluid source."""

        # Internal sources are inherit from parent class
        # Zero internal sources for unfractured domains, but we keep it to maintain
        # good code practice
        internal_sources: pp.ad.Operator = super().fluid_source(subdomains)

        # External sources are retrieved from pp.TIME_STEP_SOLUTIONS and wrapped as an
        # AdArray.
        external_sources = pp.ad.TimeDependentDenseArray(
            name="source_flow",
            domains=self.mdg.subdomains(),
            previous_timestep=True,
        )

        # Add up contribution of internal and external sources of fluid
        fluid_sources = internal_sources + external_sources
        fluid_sources.set_name("Fluid source")

        """
        The following is a "hack" to include cross term missing from the dt operator.
        This reduces the discretization error significantly, but we don't include it
        since the purpose of this verification setup is to be used in functional
        tests to check the correct *default* behaviour of the models.

        Add cross term missing from dt operator relative to proper application of
        the product rule.

        .. code:: python

            rho = self.fluid_density(subdomains)
            phi = self.volume_integral(self.porosity(subdomains), subdomains, dim=1)
            dt_op = pp.ad.time_derivatives.dt
            dt = pp.ad.Scalar(self.time_manager.dt, name="delta_t")
            prod = dt_op(rho, dt) * dt_op(phi, dt)

        """

        return fluid_sources  # - prod

class ManuPoroMechMomentumBalance(momentum.MomentumBalanceEquations):
    """Modify momentum balance to account for time-dependent body force."""

    def body_force(self, subdomains: list[pp.Grid]) -> pp.ad.Operator:
        """Body force."""

        external_sources = pp.ad.TimeDependentDenseArray(
            name="source_mechanics",
            domains=self.mdg.subdomains(),
            previous_timestep=True,
        )

        return external_sources

class ManuPoroMechEquations(
    ManuPoroMechMassBalance,
    ManuPoroMechMomentumBalance,
):
    """Mixer class for modified poromoechanics equations."""

    def set_equations(self):
        """Set the equations for the modified poromechanics problem.

        Call both parent classes' `set_equations` methods.

        """
        ManuPoroMechMassBalance.set_equations(self)
        ManuPoroMechMomentumBalance.set_equations(self)

# -----> Solution strategy
class ManuPoroMechSolutionStrategy2d(poromechanics.SolutionStrategyPoromechanics):
    """Solution strategy for the verification setup."""

    exact_sol: ManuPoroMechExactSolution2d
    """Exact solution object."""

    fluid: pp.FluidConstants
    """Object containing the fluid constants."""

    plot_results: Callable
    """Method for plotting results. Usually provided by the mixin class
    :class:`SetupUtilities`.

    """

    results: list[ManuPoroMechSaveData]
    """List of SaveData objects."""

    def __init__(self, params: dict):
        """Constructor for the class."""
        super().__init__(params)

        self.exact_sol: ManuPoroMechExactSolution2d
        """Exact solution object."""

        self.results: list[ManuPoroMechSaveData] = []
        """Results object that stores exact and approximated solutions and errors."""

        self.flux_variable: str = "darcy_flux"
        """Keyword to access the Darcy fluxes."""

        self.stress_variable: str = "poroelastic_force"
        """Keyword to access the poroelastic force."""

    def set_materials(self):
        """Set material parameters."""
        super().set_materials()

        # Instantiate exact solution object after materials have been set
        self.exact_sol = ManuPoroMechExactSolution2d(self)

    def before_nonlinear_loop(self) -> None:
        """Update values of external sources."""
        super().before_nonlinear_loop()

        sd = self.mdg.subdomains()[0]
        data = self.mdg.subdomain_data(sd)
        t = self.time_manager.time

        # Mechanics source
        mech_source = self.exact_sol.mechanics_source(sd=sd, time=t)
        pp.set_solution_values(
            name="source_mechanics", values=mech_source, data=data, time_step_index=0
        )

        # Flow source
        flow_source = self.exact_sol.flow_source(sd=sd, time=t)
        pp.set_solution_values(
            name="source_flow", values=flow_source, data=data, time_step_index=0
        )

    def after_simulation(self) -> None:
        """Method to be called after the simulation has finished."""
        if self.params.get("plot_results", False):
            self.plot_results()

# -----> Mixer class
class ManuPoroMechSetup2d(  # type: ignore[misc]
    UnitSquareGrid,
    ManuPoroMechEquations,
    ManuPoroMechSolutionStrategy2d,
    ManuPoroMechUtils,
    ManuPoroMechDataSaving,
    poromechanics.Poromechanics,
):
    """
    Mixer class for the two-dimensional non-linear poromechanics verification setup.

    """


from io import StringIO
import os
import sys
import csv
from tracemalloc import start
import numpy as np
import pandas as pd
import urllib.request
import urllib.error
import json
import datetime
import re
import zipfile
import shutil
from preprocess_data import *

datadir = './data/raw/'
droughtDir = f'{datadir}drought/'
weatherDir = f'{datadir}weather/'
featuresDir = f'{datadir}features/'
outputDir = './data/processed/'
order = ['min', 'avg', 'max', 'precip']
months = ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul', 'aug', 'sep', 'oct', 'nov', 'dec']

weatherFileName = 'weather.csv'
droughtFileName = 'drought.csv'
countyCodesName = 'county_codes.csv'
countyCoordsName = 'county_coords.csv'
populationName = 'population.csv'
featuresName = 'features.csv'

outputFileNames = [
  weatherFileName,
  droughtFileName,
  countyCodesName,
  countyCoordsName,
  populationName,
  featuresName,
]

#
def download(url, save_path, skip_download_if_save_file_exists = False, read = True):

  delete_file = False
  download_contents = None

  # construct friendly name from path filename
  if save_path is not None:
    data_name = os.path.splitext(os.path.basename(save_path))[0]
    friendly_name = data_name.replace("-", " ") \
                            .replace("_", " ")
  else:
    friendly_name = 'data'
    delete_file = True
    save_path = 'download.temp'

  def report_hook(count, block_size, total_size):
    percent = int(count * block_size * 100 / total_size)
    print(f'downloading {friendly_name}.... {(count * block_size)/1024} KB', end='\r')

  if not skip_download_if_save_file_exists or not os.path.exists(save_path):
    print(f'downloading {friendly_name}....', end='\r')
    urllib.request.urlretrieve(url, save_path, reporthook=report_hook)
    print('')

  # read contents from file
  if read:
    with open(save_path, 'r') as f:
      download_contents = f.read()

  # delete file
  if delete_file:
    os.remove(save_path)

  return download_contents

def convert_countycodes(skip_download_if_save_file_exists):
  global allStatesCounties
  id = 1

  with open(f'{outputDir}{countyCodesName}', 'w') as w:
    # header
    w.write('id INTEGER PRIMARY KEY,county_code INTEGER,fips_code INTEGER,county_name VARCHAR(50),state VARCHAR(2),country VARCHAR(3)\n')

    for state_abbr in allStatesCounties:
      state = allStatesCounties[state_abbr]
      counties = state['Counties']
      for county in counties:

        fips_code = county['Fips']
        ncdc_code = county['Ncdc']
        state_name = state['FullName']
        county_name = county['Name']

        # prepend '01' to code, indicating county is from united states
        # add 'US' value for country column
        w.write(f'{id},01{ncdc_code},{fips_code},{county_name},{state_abbr},US\n')
        id += 1

def convert_county_coords(skip_download_if_save_file_exists):
  global allStatesCounties

  # download coordinate data
  county_boundaries = download('https://public.opendatasoft.com/explore/dataset/us-county-boundaries/download/?format=csv&timezone=America/Los_Angeles&lang=en&use_labels_for_header=true&csv_separator=%3B', f'{weatherDir}us-county-boundaries.csv', skip_download_if_save_file_exists)

  # converts county coords csv
  with open(f'{outputDir}{countyCoordsName}', 'w') as w:
    # this csv contains very large fields so
    # we must increase the field size limit to something larger
    csv.field_size_limit(0x1000000)
    reader = csv.reader(county_boundaries.split('\n'), delimiter=';')
    columns = next(reader)

    # header
    w.write('county_code INTEGER PRIMARY KEY,geo_point VARCHAR(50),geo_shape TEXT[][]\n')

    # iterate lines
    id = 1
    for row in reader:
      if len(row) > 8:
        geo_point = row[0]
        geo_shape = row[1]
        state = row[8]
        county_code = row[3]
        skip = False

        if state in allStatesCounties:
          county_code = f'{allStatesCounties[state]["StateCode"]}{county_code}'
        else:
          skip = True
          print(f'skipping county coord {row[2:]}')

        if not skip:
          # process geo shape into our db format
          shape_json = json.loads(geo_shape)
          shape_processed = '"{'
          for coord in shape_json["coordinates"][0]:
            shape_processed += f'{{""{coord[0]}"",""{coord[1]}""}},'
          shape_processed = shape_processed.strip(',') + '}"'

          # prepend '01' to code, indicating county is from united states
          w.write(f'01{county_code},"{geo_point}",{shape_processed}\n')
          id += 1

def build_weather_table(skip_download_if_save_file_exists):
    filesToStrip = ['avgtmp', 'maxtmp', 'mintmp', 'precip']
    urlPaths = ['climdiv-tmpccy', 'climdiv-tmaxcy', 'climdiv-tmincy', 'climdiv-pcpncy']
    colsPrefix = ['tmp_avg', 'tmp_max', 'tmp_min', 'precip']
    dataFiles = {}

    icols = [i for i in range(len(months) + 1)]
    dtypes = [str] + [str] * len(months)
    d = pd.DataFrame(np.vstack([icols, dtypes])).to_dict(orient='records')[1]
    dff = pd.DataFrame()

    # download weather data directory listing
    weather_directory = download('https://www1.ncdc.noaa.gov/pub/data/cirs/climdiv/', None)

    # download weather data
    for filename, url_path in zip(filesToStrip, urlPaths):
      url_path_idx = weather_directory.index(url_path)
      url_path_end_idx = weather_directory.index('"', url_path_idx)
      path = weather_directory[url_path_idx:url_path_end_idx]
      dataFiles[filename] = download(f'https://www1.ncdc.noaa.gov/pub/data/cirs/climdiv/{path}', f'{weatherDir}climdiv-{filename}.txt', skip_download_if_save_file_exists)

    for filename, prefix, i in zip(filesToStrip, colsPrefix, range(len(colsPrefix))):

        # Build column names
        cols = ['id INTEGER PRIMARY KEY']
        for m in months:
            cols.append(f'{prefix}_{m} FLOAT')

        s = re.sub('( |\t)+', ' ', dataFiles[filename])
        strio = StringIO(s)
        df = pd.read_csv(strio, delimiter=' ', header=None, index_col=False, usecols=icols, dtype=d)

        # Remove datatype field, since it's the same throughout the entirity of each file
        s = df.iloc[:,0]
        s = s.str[0:5] + s.str[7:]
        df.iloc[:,0] = s

        if i == 0:
            # Add USA Country code
            cc = pd.DataFrame(['01']*len(df), dtype=(str))
            df.iloc[:,0] = cc.iloc[:,0].str[:] + df.iloc[:,0].str[:]

            # Add columns (along with id column this first time)
            df.columns = cols
            dff = pd.DataFrame(df, columns=cols)
        else:
            # Add columns
            df.columns = cols
            # Insure id parity here!
            for v1, v2 in zip(dff.iloc[:,0], df.iloc[:,0]):
                # Don't compare country code as it hasn't been added to anything but the primary id
                if v1[2:] != v2:
                    raise RuntimeError('Invalid Data Join')

            df = df.iloc[:,1:]
            dff = dff.join(df)

        print(dff)

        # Create individual files
        #df.to_csv(f'{datadir}{filename}.csv', header=False, index=False)

    # WARNING: If you open this file in Excel without specifying the first
    # column is a string, it will remove all the first zeros in the ID column
    dff.to_csv(f'{outputDir}{weatherFileName}', index=False)
    print('Succesful merge!')

def build_drought_table(skip_download_if_save_file_exists):

    urlPaths = ['climdiv-pdsist', 'climdiv-phdist', 'climdiv-pmdist', 'climdiv-sp01st', 'climdiv-sp02st', 'climdiv-sp03st', 'climdiv-sp06st', 'climdiv-sp09st', 'climdiv-sp12st', 'climdiv-sp24st']
    dataFiles = {}

    icols = [i for i in range(len(months) + 1)]
    dtypes = [str] + [str] * len(months)
    d = pd.DataFrame(np.vstack([icols, dtypes])).to_dict(orient='records')[1]
    dff = pd.DataFrame()

    # download weather data directory listing
    weather_directory = download('https://www1.ncdc.noaa.gov/pub/data/cirs/climdiv/', None)

    # download weather data
    for url_path in urlPaths:
      url_path_idx = weather_directory.index(url_path)
      url_path_end_idx = weather_directory.index('"', url_path_idx)
      path = weather_directory[url_path_idx:url_path_end_idx]
      dataFiles[url_path] = download(f'https://www1.ncdc.noaa.gov/pub/data/cirs/climdiv/{path}', f'{droughtDir}{url_path}.txt', skip_download_if_save_file_exists)

    for i, path in enumerate(urlPaths):
        datatype = path[8:]
        cols = ['id INTEGER PRIMARY KEY']
        for m in months:
            cols.append(f'{datatype}_{m} FLOAT')

        newLines = []
        lines = dataFiles[path].split('\n')
        for line in lines:
            parts = line.split()
            if len(parts) > 0:

                # TODO: Add the years of 1895 & 1896 back in. It looks like the bad
                # data comes from the rolling averages of 12 & 24 months respectively
                # (which makes sense) - but we'd need to handle this in the ui/db and
                # not allow the user to select these two values for that date range
                if int(parts[0][0:3]) > 48 or int(parts[0][-4:]) < 1897:
                    continue
                parts[0] = parts[0][1:3] + parts[0][6:]
                newLines.append(parts)

        df = pd.DataFrame(newLines, columns=cols)

        if i == 0:
            # Add USA Country code
            cc = pd.DataFrame(['01']*len(df), dtype=(str))
            df.iloc[:,0] = cc.iloc[:,0].str[:] + df.iloc[:,0].str[:]

            # Add columns (along with id column this first time)
            df.columns = cols
            dff = pd.DataFrame(df, columns=cols)
        else:
            # Add columns
            df.columns = cols
            # Insure id parity here!
            for v1, v2 in zip(dff.iloc[:,0], df.iloc[:,0]):
                # Don't compare country code as it hasn't been added to anything but the primary id
                if v1[2:] != v2:
                    raise RuntimeError('Invalid Data Join')

            df = df.iloc[:,1:]
            dff = dff.join(df)
        print(dff)
    dff.to_csv(f'{outputDir}{droughtFileName}', index=False)
    print('Succesful merge!')

def build_population_table(skip_download_if_save_file_exists):
  global allStatesCounties

  # https://data.nber.org/census/pop/cencounts.csv 1900-1990
  # https://api.census.gov/data/2000/dec/sf1?get=P001001,NAME&for=county:*

  # get historical data
  pop_1900_to_1990 = download("https://data.nber.org/census/pop/cencounts.csv", f'{datadir}pop-1900-1990.csv', skip_download_if_save_file_exists= skip_download_if_save_file_exists)

  # generates population csv
  with open(f'{outputDir}{populationName}', 'w') as w:
    reader = csv.reader(pop_1900_to_1990.split('\n'), delimiter=',')
    columns = next(reader)
    us_total_pop = next(reader)
    id = 1

    # header
    w.write('id INTEGER PRIMARY KEY,county_code INTEGER,year INTEGER,population INTEGER\n')

    # iterate historical data
    for row in reader:
      if len(row) == 12:

        # name is state code then name
        # some rows are aggregate state data
        # valid county name is "AL Autauga County"
        # agg state name is "AL Alabama"
        name = row[11]
        state_code = name[:2]
        name = name[3:]
        fips = row[10]
        county = None
        skip = False

        if fips[2:] == '000':
          # fips code of 000 means its the state/country
          # not the county
          skip = True
        elif state_code in allStatesCounties:
          counties = allStatesCounties[state_code]["Counties"]
          filtered_counties = list(filter(lambda x: x["Fips"] == fips, counties))
          if len(filtered_counties) == 1:
            county = filtered_counties[0]
          else:
            skip = True
            print(f'skipping 1900-1990 population for {state_code} {name} {fips}')
        else:
          skip = True
          print(f'skipping 1900-1990 population for {state_code} {name} {fips}')

        if not skip:

          for i in range(10):
            # prepend '01' to code, indicating county is from united states
            value = row[i]
            if value == '.':
              value = '-1'
            w.write(f'{id},01{county["Ncdc"]},{1900 + (i * 10)},{value}\n')
            id += 1

    # iterate new data until we can't
    year = 2000
    while True:

      # try and download data
      # download will throw exception on 404
      # which happens when we reach the end of the available data
      pop_data = None
      try:
        pop_data_json = download(f"https://api.census.gov/data/{year}/dec/sf1?get=P001001,NAME&for=county:*", f'{datadir}pop-{year}.csv', skip_download_if_save_file_exists= skip_download_if_save_file_exists)
        pop_data = json.loads(pop_data_json)
      except:
        break

      # first row is header
      for row in pop_data[1:]:
        population = row[0]
        name = row[1] # "county name, state name"
        state_id = row[2]
        county_fips = row[3] # just county fips, no state id
        fips = f'{state_id}{county_fips}'
        state_name = name.split(',')[1].strip()
        skip = False
        county = None

        state_code = next((s for s in allStatesCounties if allStatesCounties[s]["FullName"] == state_name), None)
        if state_code is None:
          skip = True
          print(f'skipping {year} population for {name} {state_id} {county_fips}')
        else:
          state = allStatesCounties[state_code]
          county = next((c for c in state["Counties"] if c["Fips"] == fips), None)
          if county is None:
            skip = True
            print(f'skipping {year} population for {name} {state_id} {county_fips}')

        if not skip:

          # prepend '01' to code, indicating county is from united states
          w.write(f'{id},01{county["Ncdc"]},{year},{population}\n')
          id += 1

      year += 10

def build_features_table(skip_download_if_save_file_exists):
  # https://geonames.usgs.gov/docs/stategaz/NationalFile.zip

  # get dataset
  download("https://geonames.usgs.gov/docs/stategaz/NationalFile.zip", f'{datadir}features.zip', read= False, skip_download_if_save_file_exists= skip_download_if_save_file_exists)

  # delete existing zip extraction
  if os.path.exists(featuresDir):
    shutil.rmtree(featuresDir)

  # unzip dataset
  with zipfile.ZipFile(f'{datadir}features.zip', 'r') as zip_ref:
    zip_ref.extractall(featuresDir)

  # get extracted filename
  dataset_file = os.listdir(featuresDir)[0]

  # converts features csv
  with open(f'{featuresDir}{dataset_file}', 'r', encoding='utf-8') as r:
    lines = r.readlines()

    with open(f'{outputDir}{featuresName}', 'w', encoding='utf-8') as w:
      # this csv contains very large fields so
      # we must increase the field size limit to something larger
      csv.field_size_limit(0x1000000)
      reader = csv.reader(lines, delimiter='|')
      columns = next(reader)

      # header
      w.write('id INTEGER PRIMARY KEY,county_code INTEGER,feature_type VARCHAR(50),feature_name VARCHAR(200),elevation_ft INTEGER\n')

      # iterate lines
      id = 1

      for row in reader:
        if len(row) == 20:
          name = row[1]
          name = name.replace(",","")
          feature_type = row[2]
          state_code = row[3]
          fips = f'{row[4]}{row[6]}'
          elevation = row[16]
          ncdc = None
          skip = False

          if state_code in allStatesCounties:
            ncdc = f'{allStatesCounties[state_code]["StateCode"]}{row[6]}'
          else:
            skip = True
            print(f'skipping feature {name} {feature_type} {state_code} {fips}')

          if not elevation:
            skip = True
            print(f'skipping feature {name} {feature_type} {state_code} {fips} FOR MISSING ELEVATION')

          if not skip:
            # prepend '01' to code, indicating county is from united states
            w.write(f'{id},01{ncdc},{feature_type},{name},{elevation}\n')

def has_processed_files():
  for outputFileName in outputFileNames:
    if not os.path.exists(f'{outputDir}{outputFileName}'):
      return False

  return True

def process_files(force_data_redownload = True):
  # if not forced redownload and we've already processed everything
  # exit
  if not force_data_redownload:
    if has_processed_files():
      return

  # process county codes and test the output
  try:
    build_drought_table(not force_data_redownload)
  except Exception as error:
    print(error)

  try:
    build_weather_table(not force_data_redownload)
  except Exception as error:
    print(error)

  try:
    convert_countycodes(not force_data_redownload)
  except Exception as error:
    print(error)

  try:
    convert_county_coords(not force_data_redownload)
  except Exception as error:
    print(error)

  try:
    build_population_table(not force_data_redownload)
  except Exception as error:
    print(error)

  try:
    build_features_table(not force_data_redownload)
  except Exception as error:
    print(error)

def create_working_directory():
  if not os.path.exists(outputDir):
    os.makedirs(outputDir)
  if not os.path.exists(droughtDir):
    os.makedirs(droughtDir)
  if not os.path.exists(weatherDir):
    os.makedirs(weatherDir)

if __name__ == '__main__':
  create_working_directory()
  process_files()


import pandas as pd
from copy import deepcopy
import json
from Kmeans2 import KMeans

class Method1:
    def __init__(self):
        self.str_file = '../customer_goods_info.json'
        self.customer_goods_info = None
        with open(self.str_file, 'r') as f:
            print("Load str file from {}".format(self.str_file))
            self.customer_goods_info = json.load(f)

    # 得到在n级品类结构汇总的所有商品pluno
    def get_pluno_list(self, n):
        df = pd.read_csv('../trade_new.csv')
        # 取出所需要的信息并根据pluno排序
        useful_info = df[['pluno', 'vipno', 'amt']].sort_values('pluno')
        # 取出所有的pluno列表（去重）
        pluno = useful_info.drop_duplicates(subset=['pluno'], keep='first', inplace=False)['pluno']
        pluno = [(int(item/10**(7-n))) for item in pluno]
        pluno = list(set(pluno))
        return pluno

    # 将所有商品在第n级品类结构汇总
    def n_classfication(self, n):
        data = {}
        for key in self.customer_goods_info:
            data[key] = {}
            data[key]['pluno'] = [int((int(item))/10**(7-n)) for item in self.customer_goods_info[key]['pluno']]
            data[key]['amt'] = [float(item) for item in self.customer_goods_info[key]['amt']]

        # last_data = {vipno:{pluno1:amt1, pluno2:amt2...}
        last_data = deepcopy(data)
        for index in data:
            temp_amt = 0
            temp_result = {}
            # temp_result['pluno'] = []
            # temp_result['amt'] = []
            current_pluno = data[index]['pluno'][0]
            for index2, item in enumerate(data[index]['pluno']):
                if current_pluno == item:
                    temp_amt += data[index]['amt'][index2]
                else:
                    temp_result[item] = temp_amt
                    current_pluno = item
                    temp_amt = data[index]['amt'][index2]
                if index2 == len(data[index]['pluno']) - 1:
                    temp_result[item] = temp_amt
                    break
            last_data[index] = temp_result

        return last_data

    # 计算jaccard值
    def cpmpute_jaccard(self, c1, c2):

        same_pluno_list = [x for x in c1.keys() if x in c2.keys()]
        diff_pluno_list = [y for y in (list(c1.keys()) + list(c2.keys())) if y not in same_pluno_list]

        same_list_max = []
        same_list_min = []
        different_list = []

        for item in same_pluno_list:
            if c1[item] >= c2[item]:
                same_list_max.append(c1[item])
                same_list_min.append(c2[item])
            else:
                same_list_max.append(c2[item])
                same_list_min.append(c1[item])
        for item in diff_pluno_list:
            if item in c1.keys():
                different_list.append(c1[item])
            else:
                different_list.append(c2[item])
        return sum(same_list_min) / (sum(same_list_max, sum(different_list)))

    # 得到所有jaccard值
    def get_jaccard(self, lev_n_classification_info):
        jaccard_sim = {}
        for index1 in lev_n_classification_info:
            for index2 in list(lev_n_classification_info.keys())[list(lev_n_classification_info.keys()).index(index1):]:
                if index1 == index2:
                    continue
                c1 = lev_n_classification_info[index1]
                c2 = lev_n_classification_info[index2]
                jaccard_sim[index1+'|%|'+index2] = self.cpmpute_jaccard(c1, c2)
        return jaccard_sim

    def coordinate(self, pluno_list, lev_n_classification_info):
        codi = {}
        for index in lev_n_classification_info:
            codi[index] = []
            for item in pluno_list:
                if item not in list(lev_n_classification_info[index].keys()):
                    codi[index].append(0)
                else:
                    codi[index].append(lev_n_classification_info[index][item])
        return codi

    def main(self):
        classification_info = {}
        pluno_list = {}
        # 1到4级品类结构
        for i in range(1, 5):
            classification_info[i] = self.n_classfication(i)
            pluno_list[i] = self.get_pluno_list(i)

        jaccard_sim_sum = {}
        jaccard_sim_avg = {}
        i = 0
        # 求平均jaccard
        for index in classification_info:
            if i == 0:
                jaccard_sim_sum = self.get_jaccard(classification_info[index])
            jaccard_sim = self.get_jaccard(classification_info[index])
            for index2 in jaccard_sim_sum:
                jaccard_sim_sum[index2] += jaccard_sim[index2]
            i += 1
        for index2 in jaccard_sim_sum:
            jaccard_sim_avg[index2] = jaccard_sim_sum[index2]/len(list(classification_info.keys()))

        # 求各个用户点坐标
        codi = {}
        for index in classification_info:
            codi[index] = {}
            codi[index] = self.coordinate(pluno_list[index], classification_info[index])

        sc_list = {}
        i = 20
        while i <= 30:
            print("---------------------")
            print("n_clusters: " + str(i))
            kmeans = KMeans(i, 50, codi, jaccard_sim_avg)
            result = kmeans.kMeans()
            sc = kmeans.silhouette_score(result['cluster'])
            cp = kmeans.compactness(result['cluster'])
            print("sc: " + str(sc))
            print("cp: " + str(cp))
            print("---------------------")
            sc_list[i] = sc
            i += 1

# with open("fourth_classification.json", "w") as f:
#     f.write(json.dumps(last_data, ensure_ascii=False, indent=4, separators=(',', ':')))

if __name__ == "__main__":
    method1 = Method1()
    method1.main()

    print(1)


#!/usr/local/bin/python3 -u

# Useful blog: https://stackoverflow.com/questions/56736043/extract-building-edges-from-map-image-using-python
# import the necessary packages
import argparse
import imutils
import cv2
import sys
import numpy as np
import json
import urllib.request
from urllib.error import HTTPError

def findCanvasId(mainfest, image_id):
    if isinstance(manifest['sequences'], list):
        sequences = manifest['sequences']
    else:
        sequences = [ manifest['sequences'] ]
    for sequence in sequences:
        if isinstance(sequence['canvases'], list):
            canvases = sequence['canvases']
        else:
            canvases = [ sequence['canvases'] ]
        for canvas in canvases:
            if searchForType(canvas, '@id', image_id) or searchForType(canvas, '@id', image_id[:-1]):
                return canvas['@id']
            else:
                print ('Not found {} or {} in {}'.format(image_id, image_id[:-1], canvas['@id']))
    return None

def searchForType(node, key, value):
    if isinstance(node, dict):
        if key in node and node[key] == value:
            print ('FOUND {} = {}'.format(key, value))
            return node
        else:
            for mapKey in node:
                if isinstance(node[mapKey], list) or isinstance(node[mapKey], dict):
                    response = searchForType(node[mapKey], key, value)
                    if response:
                        return response

    if isinstance(node, list):
        for element in node:
            response = searchForType(element, key, value)
            if response:
                return response
    return None

def getJson(url):
    count = 0
    while count < 3:
        try:
            response = urllib.request.urlopen(url)
            return json.loads(response.read().decode())
        except HTTPError as error:
            count += 1

    raise error

# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True, help="IIIF image ID")
ap.add_argument("-m", "--manifest", required=True, help="Manifest which contains this image")
ap.add_argument("-d", "--demo", required=False, help="Show image segmentation, useful for demos. Press space to cycle through", action="store_true")
ap.add_argument("-o", "--output", required=False, help="Filename for generated annotationlist", default="annotations.json")
ap.add_argument("-e", "--example", required=False, help="Example image showing areas found", default="demo.jpg")
ap.add_argument("-min", "--min", required=False, help="Minimum threshold for area of shape between 0 and 1", default=0.001)
ap.add_argument("-max", "--max", required=False, help="Maximum threshold for area of shape between 0 and 1", default=0.4)
ap.add_argument("-s", "--size", required=False, help="Size of IIIF image to download. Use IIIF size syntax e.g. 2000,", default='2000,')
args = vars(ap.parse_args())

iiif_id = args["image"]
demo=False
if "demo" in args and args['demo']:
    demo=True
    print ('demo: -{}-'.format(args['demo']))
if iiif_id[-1] != '/':
    iiif_id += '/'

print ('Getting URL {}'.format(args["manifest"]))
manifest = getJson(args["manifest"])

print ('Getting URL {}info.json'.format(iiif_id))
infoJson = getJson("{}info.json".format(iiif_id))

canvas_id = findCanvasId(manifest, iiif_id)
if not canvas_id:
    print ('Failed to find {} in {}'.format(iiif_id, args["manifest"]))
    sys.exit(-1)

if 'llgc.org.uk' in iiif_id and args["size"] == '2000,':
    url = '{}full/pct:40/0/default.jpg'.format(iiif_id)
else:
    url = '{}full/{}/0/default.jpg'.format(iiif_id, args["size"])

TMP_IMG_NAME = 'download.jpg'
print ('Getting {}'.format(url))
count = 0
completed=False
while count < 3:
    try:
        urllib.request.urlretrieve(url, TMP_IMG_NAME)
        completed=True
        break
    except HTTPError as error:
        count += 1

if not completed:
    raise error

# load the image, convert it to grayscale, blur it slightly,
# and threshold it
image = cv2.imread(TMP_IMG_NAME)
height, width, channels = image.shape
size_ratio = infoJson['width'] / width

print ('Ratio {}'.format(size_ratio))
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
#for i in range(-10,20):
#    #thresh = cv2.threshold(blurred, i, 255, cv2.THRESH_BINARY)[1]
#    print ('trying {}'.format(i))
#    thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,i)
#    cv2.imwrite('{}_ad_threshed.jpg'.format(i), thresh)
#sys.exit(-1)

thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2)
#thresh = cv2.threshold(blurred, 180, 255, cv2.THRESH_BINARY)[1]
canny = thresh
#canny = cv2.Canny(thresh, 50, 0, 6)
#kernel = np.ones((3,3), dtype=np.uint8)
#closed = cv2.morphologyEx(thresh, cv2.MORPH_DILATE,kernel)

#cv2.imwrite('both_threshed.jpg', thresh)

#cv2.imshow("Thresh", thresh)
#cv2.waitKey(0)
if demo:
    cv2.imshow("Image", canny)
    cv2.waitKey(0)

# find contours in the thresholded image
cnts = cv2.findContours(canny.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
cnts = imutils.grab_contours(cnts)

print ('Total contors found {}'.format(len(cnts)))
max_area = width * height
largerCnts = []
for c in cnts:
    bounding = cv2.boundingRect(c)
    contorArea = bounding[2] * bounding[3]
    if contorArea / max_area > args["min"] and contorArea / max_area < args["max"]:
        largerCnts.append(c)

cnts = largerCnts
print ('Contors within size bounds found {}'.format(len(cnts)))
annos = []
# loop over the contours
count = 1
for c in cnts:
	# compute the center of the contour
    M = cv2.moments(c)
    if M["m00"] != 0:
        #demoImage = thresh
        demoImage = image
        cX = int(M["m10"] / M["m00"])
        cY = int(M["m01"] / M["m00"])
        bounding = cv2.boundingRect(c)
        contorArea = bounding[2] * bounding[3]
        #print ('Area {} Proportion: {}, area {} maxArea: {}'.format(contorArea, ((contorArea / max_area) * 100), contorArea, max_area))
        #print ('Bounding {} bounding area {}'.format(bounding, bounding[2] * bounding[3]))
        #print (c[0][0][0])
        #print (type(c[0]))
        #print (type(c[0][0]))
        #print (type(c[0][0][0]))
        #print (len(c[0]))

        # draw the contour and center of the shape on the image
        cv2.drawContours(demoImage, [c], -1, (0, 255, 0), 2)
        cv2.circle(demoImage, (cX, cY), 7, (255, 255, 255), -1)
        cv2.putText(demoImage, "center", (cX - 20, cY - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

        # show the image
        if demo:
            cv2.imshow("Image", demoImage)
            cv2.waitKey(0)
        points = []
        path = 'M'
        for pointList in c:
            for point in pointList:
               # print (point)
                points.append((point[0].item() * size_ratio, point[1].item() * size_ratio))
                path += '{},{} '.format(point[0].item() * size_ratio, point[1].item() * size_ratio)

        path = path[:-1] + 'Z'
        #dwg = svgwrite.Drawing('test.svg', profile='tiny', size=(width * size_ratio, height * size_ratio))
        #polygon = svgwrite.shapes.Polyline(points=points)
        #dwg.add(polygon)
        #dwg.save()

        annos.append({
            "@context": "http://iiif.io/api/presentation/2/context.json",
            "@type": "oa:Annotation",
            "motivation" : [ "oa:commenting" ],
            "resource": [
                {
                    "@type" : "dctypes:Text",
                    "format" : "text/html",
                    "chars" : "
Shape {}
".format(count)
                }
            ],
            "on" :  {
                "@type" : "oa:SpecificResource",
                "within" : {
                  "@id" : "{}".format(manifest['@id']),
                  "@type" : "sc:Manifest"
                },
                "selector" : {
                  "@type" : "oa:Choice",
                  "default" : {
                    "@type" : "oa:FragmentSelector",
                    "value" : "xywh={},{},{},{}".format(bounding[0], bounding[1], bounding[2], bounding[3])
                  },
                  "item" : {
                    "@type" : "oa:SvgSelector",
                    "value" : "<svg xmlns='http://www.w3.org/2000/svg'><path xmlns=\"http://www.w3.org/2000/svg\" d=\"{}\"  id=\"rough_path_c8a67c00-247c-453a-bf83-e37d88be598b\" fill-opacity=\"0\" fill=\"#00bfff\" fill-rule=\"nonzero\" stroke=\"#00bfff\" stroke-width=\"1\" stroke-linecap=\"butt\" stroke-linejoin=\"miter\" stroke-miterlimit=\"10\" stroke-dasharray=\"\" stroke-dashoffset=\"0\" font-family=\"none\" font-weight=\"none\" font-size=\"none\" text-anchor=\"none\" style=\"mix-blend-mode: normal\"/></svg>".format(path)
                  }
                },
                "full" : "{}".format(canvas_id)
              }

        })
        count += 1

       # break;

    # <svg xmlns='http://www.w3.org/2000/svg'><path xmlns=\"http://www.w3.org/2000/svg\" d=\"M5286.33477,2466.51724l166.73851,84.88506l6.06322,118.23276l-100.0431,33.3477l-130.3592,-87.91667l57.60057,-151.58046z\" data-paper-data=\"{"strokeWidth":1,"editable":true,"deleteIcon":null,"annotation":null}\" id=\"rough_path_c8a67c00-247c-453a-bf83-e37d88be598b\" fill-opacity=\"0\" fill=\"#00bfff\" fill-rule=\"nonzero\" stroke=\"#00bfff\" stroke-width=\"1\" stroke-linecap=\"butt\" stroke-linejoin=\"miter\" stroke-miterlimit=\"10\" stroke-dasharray=\"\" stroke-dashoffset=\"0\" font-family=\"none\" font-weight=\"none\" font-size=\"none\" text-anchor=\"none\" style=\"mix-blend-mode: normal\"/></svg>

print ('Saving demo image to {}'.format(args["example"]))
cv2.imwrite(args["example"], demoImage)

annoList = {
    "@context": "http://iiif.io/api/presentation/2/context.json",
    "@id": "http://localhost:8888/examples/anno_list.json",
    "@type": "sc:AnnotationList",
    "resources": annos
}

output_filename = args["output"]
with open(output_filename, 'w') as outfile:
    json.dump(annoList, outfile, indent=4)
print ('Saved annotation list to {}'.format(output_filename))


import random

import pygame

pygame.init()
vec = pygame.math.Vector2

class Item(pygame.sprite.Sprite):
    def __init__(self, attributes, inventory):
        super().__init__()
        self.attributes = attributes
        self.inventory = inventory

        self.is_item = True

        self.surf = pygame.Surface((20, 20))
        self.surf.fill((20, 120, 30))
        self.rect = self.surf.get_rect(topleft=(random.randint(100, 500), random.randint(100, 500)))

        self.overworld_pos = self.rect.topleft

    def picked_up(self):
        if not self.inventory.inventory[len(self.inventory.inventory) - 1]:
            self.inventory.pick_up(self)
            self.kill()

    def update(self):
        pass

class Weapon(Item):
    def __init__(self, attributes, inventory):
        super().__init__(attributes, inventory)


# Testing if we can read from Wikibase

import pywikibot
from SPARQLWrapper import SPARQLWrapper, JSON

wikibase = pywikibot.Site("my", "my")
wikibase_repo = wikibase.data_repository()
wikibase_repo.login()

item = pywikibot.ItemPage(wikibase_repo, "Q1")
item.get()
print(item.labels['en']) # should print "Semantic Interoperability Community"

fr_label = "Communauté d'interopérabilité sémantique"
new_label_data = {'fr': fr_label}
item.editLabels(new_label_data, summary=u'Added French label')


import sys
import cv2
import os

# encoder (for mp4)
fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
# output file name, encoder , fps, size (fit to image size)
num = 0
while True:
    filename = f'minmax_vel{num}.mp4'
    if os.path.exists(os.path.join(os.getcwd(),filename)):
        num += 1
    else:
        break
video = cv2.VideoWriter(os.path.join(os.getcwd(),filename), fourcc, 50, (1920,1080))
print('video path -->', os.path.join(os.getcwd(),filename))

if not video.isOpened():
    print('cannot open')
    sys.exit()

pic_num = 3000
pic_file = '0'
pic_path = os.path.join(os.getcwd(),pic_file)
for i in range(0, pic_num):
    img = cv2.imread(os.path.join(pic_path,f'{i}.png'))
    if img is None:
        print('cannot read')
        break

    # add
    video.write(img)
print('generating videos....')
video.release()


import time
import math
import asyncio

import aiohttp
import requests

from .errors import ApiError

URL_TEMPLATE = "https://qrng.anu.edu.au/wp-content/plugins/colours-plugin/get_block_binary.php?_={tsm}"
BLOCK_SIZE = 128

class AsyncClientFree:
    async def _get_random_block(self, tsm):
        url = URL_TEMPLATE.format(tsm=tsm)

        async with aiohttp.ClientSession() as session:
            async with session.get(url) as resp:
                if resp.status == 200:
                    data = await resp.text()
                    return int(data, 2).to_bytes(BLOCK_SIZE, byteorder='big')
                else:
                    raise ApiError(await resp.text())

    async def get_random_bytes(self, size):
        ts = time.time()
        tsm = int(ts * 1000)
        coro_list = [
            self._get_random_block(tsm + i)
            for i in range(math.ceil(size / BLOCK_SIZE))
        ]
        block_list = await asyncio.gather(*coro_list)
        return b''.join(map(bytes, block_list))[:size]

class SyncClientFree:
    def _get_random_block(self, tsm):
        url = URL_TEMPLATE.format(tsm=tsm)

        with requests.get(url) as resp:
            if resp.status_code == 200:
                return int(resp.text, 2).to_bytes(BLOCK_SIZE, byteorder='big')
            else:
                raise ApiError(resp.text)

    def get_random_bytes(self, size):
        ts = time.time()
        tsm = int(ts * 1000)
        block_list = []
        for i in range(math.ceil(size / BLOCK_SIZE)):
            block = self._get_random_block(tsm + i)
            block_list.append(block)
        return b''.join(map(bytes, block_list))[:size]


import numpy as np
from numba import njit,float64
from numba.pycc import CC

cc = CC('model_precompiled_functions')
cc.verbose = True

## Inner product function
@njit(float64(float64[:],float64[:]))
def innerproduct(v1,v2):
    x=v1[0]*v2[0]+v1[1]*v2[1]+v1[2]*v2[2]
    return x
## Inner matrix product function
@njit(float64(float64[:,:],float64[:,:]))
def inner_matrix_product(A1,A2):
    x=0.0
    for i in range(3):
        for j in range(3):
            x+=A1[i,j]*A2[i,j]
    return x
## Centrifugal forces and stiffness matrix
#
# Inputs:
#  ndofs             int32
#  jcol_nonzero_irow int32[:,:]
#  ijsym             int32[:,:]
#  m_drcg_dqi        float64[:,:]
#  m_ddrcg_dqidqj    float64[:,:]
#  Abase_i           float64[:,:,:]
#  Abase1_ij         float64[:,:,:]
#  Abase2_ij         float64[:,:,:]
#  R0Tddr0           float64[:]
#  R0TddR0           float64[:,:]
#  R0TdR0            float64[:,:]
#
# Outputs
#  Fc          float64[:]
#  Gc          float64[:,:]
#  Kc          float64[:,:]
#
#
@cc.export('compute_local_centrifugal_forces_and_matrix','(i4,i4[:,:],i4[:,:],f8[:,:],f8[:,:],f8[:,:,:],f8[:,:,:],f8[:,:,:],f8[:],f8[:,:],f8[:,:])')
def compute_local_centrifugal_forces_and_matrix(ndofs,jcol_nonzero_irow,ijsym,m_drcg_dqi,m_ddrcg_dqidqj,Abase_i,Abase1_ij,Abase2_ij,R0Tddr0,R0TddR0,R0TdR0):
    ddR0TR0=R0TddR0.T
    Fc=np.zeros(ndofs)
    Gc=np.zeros((ndofs,ndofs))
    Kc=np.zeros((ndofs,ndofs))
    for i in range(ndofs):
        Fc[i]= innerproduct(R0Tddr0,m_drcg_dqi[:,i]) \
              +inner_matrix_product(ddR0TR0,Abase_i[:,:,i])
        for j in range(jcol_nonzero_irow[0,i],jcol_nonzero_irow[1,i]):
            # Upper triangle
            Gc[i,j]= inner_matrix_product(R0TdR0,Abase1_ij[:,:,ijsym[i,j]])
            Kc[i,j]= innerproduct(R0Tddr0,m_ddrcg_dqidqj[:,ijsym[i,j]]) \
                    +inner_matrix_product(R0TddR0,Abase1_ij[:,:,ijsym[i,j]]) \
                    +inner_matrix_product(ddR0TR0,Abase2_ij[:,:,ijsym[i,j]])
            # Lower triangle
            if j!=i:
                Gc[i,j]= inner_matrix_product(R0TdR0,Abase1_ij[:,:,ijsym[i,j]].T)
                Kc[j,i]= innerproduct(R0Tddr0,m_ddrcg_dqidqj[:,ijsym[i,j]]) \
                        +inner_matrix_product(R0TddR0,Abase1_ij[:,:,ijsym[i,j]].T) \
                        +inner_matrix_product(ddR0TR0,Abase2_ij[:,:,ijsym[i,j]])

    return Fc,Gc,Kc

#===============================================================================================
# Compilation
#===============================================================================================
if __name__ == "__main__":
    cc.compile()


import matplotlib.pyplot as plt
import os
import funcs

colors = list('rgbcm')

# getting all commits
all_commits = os.popen('git log --pretty=oneline --all').readlines()
all_commits = [commit.split()[0] for commit in all_commits]

# creating dictionary that assigns a uniqe index for each commit
all_commits_dict = funcs.create_commits_dict(all_commits)

# getting master branch commits
master_commits = funcs.get_master_branch_commits()

# getting the name of all exsiting branchs
branches = funcs.get_branches()

# assgning a uniqe index for each branch
branches_dict = funcs.assign_axis_to_branches(branches)

# creating a dictionary that adds every commit to each related branch
# and plots them at the same time
commit_branch_rel = {}
for branch in branches:
    commits = funcs.get_commits_of_branch(branch)
    commit_branch_rel[branch] = commits
    for commit in commits:
        plt.scatter(all_commits_dict[commit], branches_dict[branch], color=colors[0],zorder=2,s=500)
    first_color = colors[0]
    colors.remove(colors[0])
    colors.append(first_color)
commit_branch_rel['master'] = master_commits

# drawing master branch commits
for commit in master_commits:
    plt.scatter(all_commits_dict[commit], branches_dict['master'], color='y',zorder=2, s=500)

# draw lines between related commits
for commit in all_commits:
    parents = funcs.get_parents_of_commit(commit)
    for parent in parents:
        parent_branch = funcs.get_branch_name_of_commit_from_dictionary(parent, commit_branch_rel)
        commit_branch = funcs.get_branch_name_of_commit_from_dictionary(commit, commit_branch_rel)
        plt.plot([all_commits_dict[parent],all_commits_dict[commit]], [branches_dict[parent_branch],branches_dict[commit_branch]], color='black',
                 zorder=1,
                 linewidth=1)

plt.show()


import functools

def mycmp(x, y):
    return int(x) - int(y)

def resolve(l):
    l = list(map(int, l))
    return max(l[0] * l[1] * l[-1], l[-1] * l[-2] * l[-3])

num = int(input())
for i in range(num):
    people = int(input())
    l = input().split(" ")
    print(resolve(sorted(l, key=functools.cmp_to_key(mycmp))))


#servidor de echo: lado cliente
import rpyc

HOST = 'localhost' # maquina onde esta o servidor
PORT = 10020       # porta que o servidor esta escutando

def iniciaCliente():
	'''Conecta-se ao servidor.
	Saida: retorna a conexao criada.'''
	conn = rpyc.connect(HOST, PORT)

	print(type(conn.root)) # mostra que conn.root eh um stub de cliente
	print(conn.root.get_service_name()) # exibe o nome da classe (servico) oferecido

	return conn

def fazRequisicoes(conn):
	# le as mensagens do usuario ate ele digitar 'fim'
	while True:
		msg = input("\nO que deseja fazer?\nOpções:\n > Adicionar palavra/significado: Digite 'Add <palavra> <significado>'\n > Listar significados: Digite 'List <palavra>'\n > Deletar palavra: Digite 'Del <palavra>'\n > Fechar conexão: Digite 'Fim'\n\n")
		if msg == 'Fim': break

		# envia a mensagem do usuario para o servidor
		ret = conn.root.exposed_atendeRequisicoes(msg)

		print(ret)

	# encerra a conexao
	conn.close()

def main():
	#inicia o cliente
	conn = iniciaCliente()
	#interage com o servidor ate encerrar
	fazRequisicoes(conn)

main()


import pika
import time
import random

class SensorPresenca:
    def __init__(self) -> None:
        pass

    def verificaPresenca(self):
        presenca = random.randint(0, 1)
        return str(presenca)

# Configurar a conexão RabbitMQ
connection = pika.BlockingConnection(pika.ConnectionParameters('localhost'))
channel = connection.channel()

# Declarar uma fila chamada 'presenca_queue'
channel.queue_declare(queue='presenca_queue')

while True:
    sensor = SensorPresenca()
    presenca = sensor.verificaPresenca()
    print('Presença:', presenca)

    # Publicar a mensagem no RabbitMQ
    channel.basic_publish(exchange='', routing_key='presenca_queue', body=presenca)
    print(f'Mensagem {presenca} publicada na fila')
    time.sleep(5)

# Fechar a conexão com o RabbitMQ
connection.close()


"""
This will connect the Sickchill platform to Homeassistant, showing stats and switches from Sickchill.
"""
import logging
import time
import re
from datetime import timedelta
import requests
import urllib

#import newznab
from newznab import wrapper

COMPONENT_NAME = "sabnzbd"

class sabnzbd:
	def __init__(self,config,_LOGGER):
		"""Set up the component."""
		self._LOGGER = _LOGGER
		self.category = config['sab_category']
		self.sab_url = config['sab_host']
		self.sab_api_key = config['sab_api_key']

		_LOGGER.warning(' %s is starting', COMPONENT_NAME)

	def addnzb(self,url,name=None):
		'''
		Send a NZB url to SABNZBD

		:param url: The NZBSearchResult object to send to SAB
		'''

		self._LOGGER.debug("Adding url to SABNZBD: %s",url)

		# set up a dict with the URL params in it
		params = {'output': 'json'}
		params['apikey'] = self.sab_api_key

		# Set the sab catgory for sickchill
		params['cat'] = self.category

		self._LOGGER.info('Sending NZB to SABnzbd')
		submiturl = urllib.parse.urljoin(self.sab_url, 'api')

		params['mode'] = 'addurl'
		params['name'] = url

		if name is not None:
			params['nzbname'] = name

		result = requests.post(url=submiturl,data=params)

		self._LOGGER.debug("REturned result form sab: %s",result.status_code)


'''
977. Squares of a Sorted Array
Easy

Given an integer array nums sorted in non-decreasing order, return an array
of the squares of each number sorted in non-decreasing order.

Example 1:
Input: nums = [-4,-1,0,3,10]
Output: [0,1,9,16,100]
Explanation: After squaring, the array becomes [16,1,0,9,100].
After sorting, it becomes [0,1,9,16,100].

Example 2:
Input: nums = [-7,-3,2,3,11]
Output: [4,9,9,49,121]

Constraints:
1 <= nums.length <= 10**4
-10**4 <= nums[i] <= 10**4
nums is sorted in non-decreasing order.

Follow up: Squaring each element and sorting the new array is very trivial,
could you find an O(n) solution using a different approach?
'''

from typing import List

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        if nums[0] >= 0 or len(nums) == 1:
            return [x**2 for x in nums]

        left = 0
        right = len(nums) - 1
        result = []
        while left != right:
            if abs(nums[left]) >= abs(nums[right]):
                result.append(nums[left] ** 2)
                left += 1
            else:
                result.append(nums[right] ** 2)
                right -= 1

        result.append(nums[left] ** 2)
        return result[::-1]

sol = Solution()

assert sol.sortedSquares([0, 3, 10]) == [0, 9, 100]
assert sol.sortedSquares([-4, -1, 0, 3, 10]) == [0, 1, 9, 16, 100]
assert sol.sortedSquares([-7, -3, 2, 3, 11]) == [4, 9, 9, 49, 121]


# This file is part of the Printrun suite.
#
# Printrun is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Printrun is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Printrun.  If not, see <http://www.gnu.org/licenses/>.

from queue import Queue
from collections import deque
import numpy
import wx
import time
from . import gcoder
from .injectgcode import injector, injector_edit
from printrun.gui.viz import BaseViz
from .gui.widgets import get_space

from .utils import imagefile, install_locale, get_home_pos
install_locale('pronterface')

class GvizBaseFrame(wx.Frame):

    def create_base_ui(self):
        self.CreateStatusBar(1)
        self.SetStatusText(_("Layer number and Z position show here when you scroll"))

        panel = wx.Panel(self, -1)
        v_sizer = wx.BoxSizer(wx.VERTICAL)
        h_sizer = wx.BoxSizer(wx.HORIZONTAL)

        self.toolbar = wx.ToolBar(panel, -1, style = wx.TB_HORIZONTAL | wx.NO_BORDER | wx.TB_HORZ_TEXT)
        self.build_toolbar()
        self.toolbar.Realize()
        v_sizer.Add(self.toolbar, 0, wx.EXPAND)
        v_sizer.Add(h_sizer, 1, wx.EXPAND)

        self.layerslider = wx.Slider(panel, style = wx.SL_VERTICAL | wx.SL_AUTOTICKS | wx.SL_LEFT | wx.SL_INVERSE)
        self.layerslider.Bind(wx.EVT_SCROLL, self.process_slider)
        panel.SetSizer(v_sizer)

        return panel, h_sizer

    def build_toolbar(self):
        self.toolbar.SetMargins(get_space('minor'), get_space('mini'))
        self.toolbar.SetToolPacking(get_space('minor'))
        self.toolbar.AddTool(1, '', wx.Image(imagefile('zoom_out.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(), _("Zoom Out [-]"))
        self.toolbar.AddTool(2, '', wx.Image(imagefile('zoom_in.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(), _("Zoom In [+]"),)
        self.toolbar.AddTool(3, _("Reset View"), wx.Image(imagefile('fit.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(),
                             shortHelp = _("Reset View"))
        self.toolbar.AddSeparator()
        self.toolbar.AddTool(4, '', wx.Image(imagefile('arrow_down.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(), _("Move Down a Layer [D]"))
        self.toolbar.AddTool(5, '', wx.Image(imagefile('arrow_up.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(), _("Move Up a Layer [U]"))
        self.toolbar.AddSeparator()
        self.toolbar.AddTool(6, 'Inject', wx.Image(imagefile('inject.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(),
                             wx.NullBitmap, shortHelp = _("Inject G-Code"), longHelp = _("Insert code at the beginning of this layer"))
        self.toolbar.AddTool(7, 'Edit', wx.Image(imagefile('edit.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(),
                             wx.NullBitmap, shortHelp = _("Edit Layer"), longHelp = _("Edit the G-Code of this layer"))
        self.toolbar.AddStretchableSpace()
        self.toolbar.AddSeparator()
        self.toolbar.AddTool(9, _('Close'), wx.Image(imagefile('fit.png'), wx.BITMAP_TYPE_PNG).ConvertToBitmap(),
                             shortHelp = _("Close Window"))

    def setlayercb(self, layer):
        self.layerslider.SetValue(layer)

    def process_slider(self, event):
        raise NotImplementedError

ID_ABOUT = 101
ID_EXIT = 110
class GvizWindow(GvizBaseFrame):
    def __init__(self, f = None, size = (550, 550), build_dimensions = [200, 200, 100, 0, 0, 0], grid = (10, 50), extrusion_width = 0.5, bgcolor = "#000000"):
        super().__init__(None, title = _("G-Code Viewer"), size = size, style = wx.DEFAULT_FRAME_STYLE)

        panel, h_sizer = self.create_base_ui()

        self.p = Gviz(panel, size = size, build_dimensions = build_dimensions, grid = grid, extrusion_width = extrusion_width, bgcolor = bgcolor, realparent = self)
        h_sizer.Add(self.p, 1, wx.EXPAND)
        h_sizer.Add(self.layerslider, 0, wx.EXPAND | wx.ALL, get_space('minor'))
        self.p.SetToolTip("Left-click to move the view, scroll to zoom and shift + scroll to change the layer.")

        self.Layout()
        size = (size[0] + self.layerslider.GetEffectiveMinSize().width,
                size[1] + self.toolbar.GetEffectiveMinSize().height)
        minsize = self.toolbar.GetEffectiveMinSize().width
        self.SetMinClientSize((minsize, minsize))
        self.SetClientSize(size)

        self.Bind(wx.EVT_TOOL, lambda x: self.p.zoom(-1, -1, 1 / 1.2), id = 1)
        self.Bind(wx.EVT_TOOL, lambda x: self.p.zoom(-1, -1, 1.2), id = 2)
        self.Bind(wx.EVT_TOOL, self.reset_view, id = 3)
        self.Bind(wx.EVT_TOOL, lambda x: self.p.layerdown(), id = 4)
        self.Bind(wx.EVT_TOOL, lambda x: self.p.layerup(), id = 5)
        self.Bind(wx.EVT_TOOL, lambda x: self.p.inject(), id = 6)
        self.Bind(wx.EVT_TOOL, lambda x: self.p.editlayer(), id = 7)
        self.Bind(wx.EVT_TOOL, self.reset_selection, id = 8)
        self.Bind(wx.EVT_TOOL, lambda x: self.Close(), id = 9)

        self.initpos = None
        self.p.Bind(wx.EVT_KEY_DOWN, self.key)
        self.Bind(wx.EVT_KEY_DOWN, self.key)
        self.p.Bind(wx.EVT_MOUSEWHEEL, self.zoom)
        self.Bind(wx.EVT_MOUSEWHEEL, self.zoom)
        self.p.Bind(wx.EVT_MOUSE_EVENTS, self.mouse)
        self.Bind(wx.EVT_MOUSE_EVENTS, self.mouse)

        if f:
            gcode = gcoder.GCode(f, get_home_pos(self.p.build_dimensions))
            self.p.addfile(gcode)

        self.CenterOnParent()

    def set_current_gline(self, gline):
        return

    def process_slider(self, event):
        self.p.layerindex = self.layerslider.GetValue()
        z = self.p.get_currentz()
        wx.CallAfter(self.SetStatusText, _("Layer %d: Z = %.03f mm") % (self.p.layerindex + 1, z), 0)
        self.p.dirty = True
        wx.CallAfter(self.p.Refresh)

    def reset_view(self, event):
        self.p.translate = [0.0, 0.0]
        self.p.scale = self.p.basescale
        self.p.zoom(0, 0, 1.0)

    def reset_selection(self, event):
        self.parent.rectangles = []
        wx.CallAfter(self.p.Refresh)

    def mouse(self, event):
        if event.ButtonUp(wx.MOUSE_BTN_LEFT) or event.ButtonUp(wx.MOUSE_BTN_RIGHT):
            if self.initpos is not None:
                self.initpos = None
        elif event.Dragging():
            e = event.GetPosition()
            if self.initpos is None:
                self.initpos = e
                self.basetrans = self.p.translate
            self.p.translate = [self.basetrans[0] + (e[0] - self.initpos[0]),
                                self.basetrans[1] + (e[1] - self.initpos[1])]
            self.p.dirty = True
            wx.CallAfter(self.p.Refresh)
        else:
            event.Skip()

    def key(self, event):
        #  Keycode definitions
        kup = [85, 315]               # Up keys
        kdo = [68, 317]               # Down Keys
        kzi = [388, 316, 61]        # Zoom In Keys
        kzo = [390, 314, 45]       # Zoom Out Keys
        x = event.GetKeyCode()
        cx, cy = self.p.translate
        if x in kup:
            self.p.layerup()
        if x in kdo:
            self.p.layerdown()
        if x in kzi:
            self.p.zoom(cx, cy, 1.2)
        if x in kzo:
            self.p.zoom(cx, cy, 1 / 1.2)

    def zoom(self, event):
        z = event.GetWheelRotation()
        if event.ShiftDown():
            if z > 0:
                self.p.layerdown()
            elif z < 0:
                self.p.layerup()
        else:
            if z > 0:
                self.p.zoom(event.GetX(), event.GetY(), 1.2)
            elif z < 0:
                self.p.zoom(event.GetX(), event.GetY(), 1 / 1.2)

class Gviz(wx.Panel, BaseViz):

    # Mark canvas as dirty when setting showall
    _showall = 0

    def _get_showall(self):
        return self._showall

    def _set_showall(self, showall):
        if showall != self._showall:
            self.dirty = True
            self._showall = showall
    showall = property(_get_showall, _set_showall)

    def __init__(self, parent, size = (200, 200), build_dimensions = [200, 200, 100, 0, 0, 0], grid = (10, 50), extrusion_width = 0.5, bgcolor = "#000000", realparent = None):
        wx.Panel.__init__(self, parent, -1)
        self.widget = self
        size = [max(1.0, x) for x in size]
        ratio = size[0] / size[1]
        self.SetMinSize((150, 150 / ratio))
        self.parent = realparent if realparent else parent
        self.size = size
        self.build_dimensions = build_dimensions
        self.grid = grid
        self.Bind(wx.EVT_PAINT, self.paint)
        self.Bind(wx.EVT_SIZE, self.resize)
        self.hilight = deque()
        self.hilightarcs = deque()
        self.hilightqueue = Queue(0)
        self.hilightarcsqueue = Queue(0)
        self.clear()
        self.filament_width = extrusion_width  # set it to 0 to disable scaling lines with zoom
        self.update_basescale()
        self.scale = self.basescale
        penwidth = max(1, int(self.filament_width * ((self.scale[0] + self.scale[1]) / 2.0)))
        self.translate = [0.0, 0.0]
        self.mainpen = wx.Pen(wx.Colour(0, 0, 0), penwidth)
        self.arcpen = wx.Pen(wx.Colour(255, 0, 0), penwidth)
        self.travelpen = wx.Pen(wx.Colour(10, 80, 80), penwidth)
        self.hlpen = wx.Pen(wx.Colour(200, 50, 50), penwidth)
        self.fades = [wx.Pen(wx.Colour(int(250 - 0.6 ** i * 100), int(250 - 0.6 ** i * 100), int(200 - 0.4 ** i * 50)), penwidth) for i in range(6)]
        self.penslist = [self.mainpen, self.arcpen, self.travelpen, self.hlpen] + self.fades
        self.bgcolor = wx.Colour()
        self.bgcolor.Set(bgcolor)
        self.blitmap = wx.Bitmap(self.GetClientSize()[0], self.GetClientSize()[1], -1)
        self.paint_overlay = None

    def inject(self):
        layer = self.layers[self.layerindex]
        injector(self.gcode, self.layerindex, layer)

    def editlayer(self):
        layer = self.layers[self.layerindex]
        injector_edit(self.gcode, self.layerindex, layer)

    def clearhilights(self):
        self.hilight.clear()
        self.hilightarcs.clear()
        while not self.hilightqueue.empty():
            self.hilightqueue.get_nowait()
        while not self.hilightarcsqueue.empty():
            self.hilightarcsqueue.get_nowait()

    def clear(self):
        self.gcode = None
        self.lastpos = [0, 0, 0, 0, 0, 0, 0]
        self.hilightpos = self.lastpos[:]
        self.lines = {}
        self.pens = {}
        self.arcs = {}
        self.arcpens = {}
        self.layers = {}
        self.layersz = []
        self.clearhilights()
        self.layerindex = 0
        self.showall = 0
        self.dirty = True
        self.partial = False
        self.painted_layers = set()
        wx.CallAfter(self.Refresh)

    def get_currentz(self):
        z = self.layersz[self.layerindex]
        z = 0. if z is None else z
        return z

    def layerup(self):
        if self.layerindex + 1 < len(self.layers):
            self.layerindex += 1
            z = self.get_currentz()
            wx.CallAfter(self.parent.SetStatusText, _("Layer %d: Z = %.03f mm") % (self.layerindex + 1, z), 0)
            self.dirty = True
            self.parent.setlayercb(self.layerindex)
            wx.CallAfter(self.Refresh)

    def layerdown(self):
        if self.layerindex > 0:
            self.layerindex -= 1
            z = self.get_currentz()
            wx.CallAfter(self.parent.SetStatusText, _("Layer %d: Z = %.03f mm") % (self.layerindex + 1, z), 0)
            self.dirty = True
            self.parent.setlayercb(self.layerindex)
            wx.CallAfter(self.Refresh)

    def setlayer(self, layer):
        if layer in self.layers:
            self.clearhilights()
            self.layerindex = self.layers[layer]
            self.dirty = True
            self.showall = 0
            wx.CallAfter(self.Refresh)

    def update_basescale(self):
        self.basescale = 2 * [min(float(self.size[0] - 1) / self.build_dimensions[0],
                                  float(self.size[1] - 1) / self.build_dimensions[1])]

    def resize(self, event):
        old_basescale = self.basescale
        width, height = self.GetClientSize()
        if width < 1 or height < 1:
            return
        self.size = (width, height)
        self.update_basescale()
        zoomratio = float(self.basescale[0]) / old_basescale[0]
        wx.CallLater(200, self.zoom, 0, 0, zoomratio)

    def zoom(self, x, y, factor):
        if x == -1 and y == -1:
            side = min(self.size)
            x = y = side / 2
        self.scale = [s * factor for s in self.scale]

        self.translate = [x - (x - self.translate[0]) * factor,
                          y - (y - self.translate[1]) * factor]
        penwidth = max(1, int(self.filament_width * ((self.scale[0] + self.scale[1]) / 2.0)))
        for pen in self.penslist:
            pen.SetWidth(penwidth)
        self.dirty = True
        wx.CallAfter(self.Refresh)

    def _line_scaler(self, x):
        return (int(self.scale[0] * x[0]),
                int(self.scale[1] * x[1]),
                int(self.scale[0] * x[2]),
                int(self.scale[1] * x[3]),)

    def _arc_scaler(self, x):
        return (self.scale[0] * x[0],
                self.scale[1] * x[1],
                self.scale[0] * x[2],
                self.scale[1] * x[3],
                self.scale[0] * x[4],
                self.scale[1] * x[5],)

    def _drawlines(self, dc, lines, pens):
        scaled_lines = [self._line_scaler(l) for l in lines]
        dc.DrawLineList(scaled_lines, pens)

    def _drawarcs(self, dc, arcs, pens):
        scaled_arcs = [self._arc_scaler(a) for a in arcs]
        dc.SetBrush(wx.TRANSPARENT_BRUSH)
        for i, arc in enumerate(scaled_arcs):
            dc.SetPen(pens[i] if isinstance(pens, numpy.ndarray) else pens)
            dc.DrawArc(*arc)

    def repaint_everything(self):
        width = self.scale[0] * self.build_dimensions[0]
        height = self.scale[1] * self.build_dimensions[1]
        self.blitmap = wx.Bitmap(int(width) + 1, int(height) + 1, -1)
        dc = wx.MemoryDC()
        dc.SelectObject(self.blitmap)
        dc.SetBackground(wx.Brush((250, 250, 200)))
        dc.Clear()
        dc.SetPen(wx.Pen(wx.Colour(180, 180, 150)))
        for grid_unit in self.grid:
            if grid_unit > 0:
                for x in range(int(self.build_dimensions[0] / grid_unit) + 1):
                    draw_x = self.scale[0] * x * grid_unit
                    dc.DrawLine(int(draw_x), 0, int(draw_x), int(height))
                for y in range(int(self.build_dimensions[1] / grid_unit) + 1):
                    draw_y = self.scale[1] * (self.build_dimensions[1] - y * grid_unit)
                    dc.DrawLine(0, int(draw_y), int(width), int(draw_y))
            dc.SetPen(wx.Pen(wx.Colour(0, 0, 0)))

        if not self.showall:
            # Draw layer gauge
            dc.SetBrush(wx.Brush((43, 144, 255)))
            dc.DrawRectangle(int(width) - 15, 0, 15, int(height))
            dc.SetBrush(wx.Brush((0, 255, 0)))
            if self.layers:
                dc.DrawRectangle(int(width) - 14, int((1.0 - (1.0 * (self.layerindex + 1)) / len(self.layers)) * height), 13, int(height) - 1)

        if self.showall:
            for i in range(len(self.layersz)):
                self.painted_layers.add(i)
                self._drawlines(dc, self.lines[i], self.pens[i])
                self._drawarcs(dc, self.arcs[i], self.arcpens[i])
            dc.SelectObject(wx.NullBitmap)
            return

        if self.layerindex < len(self.layers) and self.layerindex in self.lines:
            for layer_i in range(max(0, self.layerindex - 6), self.layerindex):
                self._drawlines(dc, self.lines[layer_i], self.fades[self.layerindex - layer_i - 1])
                self._drawarcs(dc, self.arcs[layer_i], self.fades[self.layerindex - layer_i - 1])
            self._drawlines(dc, self.lines[self.layerindex], self.pens[self.layerindex])
            self._drawarcs(dc, self.arcs[self.layerindex], self.arcpens[self.layerindex])

        self._drawlines(dc, self.hilight, self.hlpen)
        self._drawarcs(dc, self.hilightarcs, self.hlpen)

        self.paint_hilights(dc)

        dc.SelectObject(wx.NullBitmap)

    def repaint_partial(self):
        if self.showall:
            dc = wx.MemoryDC()
            dc.SelectObject(self.blitmap)
            for i in set(range(len(self.layersz))).difference(self.painted_layers):
                self.painted_layers.add(i)
                self._drawlines(dc, self.lines[i], self.pens[i])
                self._drawarcs(dc, self.arcs[i], self.arcpens[i])
            dc.SelectObject(wx.NullBitmap)

    def paint_hilights(self, dc = None):
        if self.hilightqueue.empty() and self.hilightarcsqueue.empty():
            return
        hl = []
        if not dc:
            dc = wx.MemoryDC()
            dc.SelectObject(self.blitmap)
        while not self.hilightqueue.empty():
            hl.append(self.hilightqueue.get_nowait())
        self._drawlines(dc, hl, self.hlpen)
        hlarcs = []
        while not self.hilightarcsqueue.empty():
            hlarcs.append(self.hilightarcsqueue.get_nowait())
        self._drawarcs(dc, hlarcs, self.hlpen)
        dc.SelectObject(wx.NullBitmap)

    def paint(self, event):
        if self.dirty:
            self.dirty = False
            self.partial = False
            self.repaint_everything()
        elif self.partial:
            self.partial = False
            self.repaint_partial()
        self.paint_hilights()
        dc = wx.PaintDC(self)
        dc.SetBackground(wx.Brush(self.bgcolor))
        dc.Clear()
        dc.DrawBitmap(self.blitmap, int(self.translate[0]), int(self.translate[1]))
        if self.paint_overlay:
            self.paint_overlay(dc)

    def addfile_perlayer(self, gcode, showall = False):
        self.clear()
        self.gcode = gcode
        self.showall = showall
        generator = self.add_parsed_gcodes(gcode)
        generator_output = next(generator)
        while generator_output is not None:
            yield generator_output
            generator_output = next(generator)
        max_layers = len(self.layers)
        if hasattr(self.parent, "layerslider"):
            self.parent.layerslider.SetRange(0, max_layers - 1)
            self.parent.layerslider.SetValue(0)
        yield None

    def addfile(self, gcode = None, showall = False):
        generator = self.addfile_perlayer(gcode, showall)
        while next(generator) is not None:
            continue

    def _get_movement(self, start_pos, gline):
        """Takes a start position and a gcode, and returns a 3-uple containing
        (final position, line, arc), with line and arc being None if not
        used"""
        target = start_pos[:]
        target[5] = 0.0
        target[6] = 0.0
        if gline.current_x is not None:
            target[0] = gline.current_x
        if gline.current_y is not None:
            target[1] = gline.current_y
        if gline.current_z is not None:
            target[2] = gline.current_z
        if gline.e is not None:
            if gline.relative_e:
                target[3] += gline.e
            else:
                target[3] = gline.e
        if gline.f is not None:
            target[4] = gline.f
        if gline.i is not None:
            target[5] = gline.i
        if gline.j is not None:
            target[6] = gline.j

        if gline.command in ["G2", "G3"]:
            # startpos, endpos, arc center
            arc = [self._x(start_pos[0]), self._y(start_pos[1]),
                   self._x(target[0]), self._y(target[1]),
                   self._x(start_pos[0] + target[5]), self._y(start_pos[1] + target[6])]
            if gline.command == "G2":  # clockwise, reverse endpoints
                arc[0], arc[1], arc[2], arc[3] = arc[2], arc[3], arc[0], arc[1]
            return target, None, arc

        # ["G0", "G1"]:
        line = [self._x(start_pos[0]),
                self._y(start_pos[1]),
                self._x(target[0]),
                self._y(target[1])]
        return target, line, None

    def _y(self, y):
        return self.build_dimensions[1] - (y - self.build_dimensions[4])

    def _x(self, x):
        return x - self.build_dimensions[3]

    def add_parsed_gcodes(self, gcode):
        start_time = time.time()

        layer_idx = 0
        while layer_idx < len(gcode.all_layers):
            layer = gcode.all_layers[layer_idx]
            has_move = False
            for gline in layer:
                if gline.is_move:
                    has_move = True
                    break
            if not has_move:
                yield layer_idx
                layer_idx += 1
                continue
            viz_layer = len(self.layers)
            self.lines[viz_layer] = []
            self.pens[viz_layer] = []
            self.arcs[viz_layer] = []
            self.arcpens[viz_layer] = []
            for gline in layer:
                if not gline.is_move:
                    continue

                target, line, arc = self._get_movement(self.lastpos[:], gline)

                if line is not None:
                    self.lines[viz_layer].append(line)
                    self.pens[viz_layer].append(self.mainpen if target[3] != self.lastpos[3] or gline.extruding else self.travelpen)
                elif arc is not None:
                    self.arcs[viz_layer].append(arc)
                    self.arcpens[viz_layer].append(self.arcpen)

                self.lastpos = target
            # Transform into a numpy array for memory efficiency
            self.lines[viz_layer] = numpy.asarray(self.lines[viz_layer], dtype = numpy.float32)
            self.pens[viz_layer] = numpy.asarray(self.pens[viz_layer])
            self.arcs[viz_layer] = numpy.asarray(self.arcs[viz_layer], dtype = numpy.float32)
            self.arcpens[viz_layer] = numpy.asarray(self.arcpens[viz_layer])
            # Only add layer to self.layers now to prevent the display of an
            # unfinished layer
            self.layers[layer_idx] = viz_layer
            self.layersz.append(layer.z)

            # Refresh display if more than 0.2s have passed
            if time.time() - start_time > 0.2:
                start_time = time.time()
                self.partial = True
                wx.CallAfter(self.Refresh)

            yield layer_idx
            layer_idx += 1

        self.dirty = True
        wx.CallAfter(self.Refresh)
        yield None

    def addgcodehighlight(self, gline):
        if gline.command not in ["G0", "G1", "G2", "G3"]:
            return

        target, line, arc = self._get_movement(self.hilightpos[:], gline)

        if line is not None:
            self.hilight.append(line)
            self.hilightqueue.put_nowait(line)
        elif arc is not None:
            self.hilightarcs.append(arc)
            self.hilightarcsqueue.put_nowait(arc)

        self.hilightpos = target
        wx.CallAfter(self.Refresh)

if __name__ == '__main__':
    import sys
    app = wx.App(False)
    with open(sys.argv[1], "r") as arg:
        main = GvizWindow(arg)
    main.Show()
    app.MainLoop()