from __future__ import annotations
"""
.. module:: tree
"""
import attr
import math
import numpy as np

from typing import List, Optional, TYPE_CHECKING

import pydynamo_brain.util as util
from pydynamo_brain.util import SAVE_META, Point3D

if TYPE_CHECKING:
    from .branch import Branch

@attr.s
class Point():
    """Node in the tree, a point in 3D space."""

    id: str = attr.ib(metadata=SAVE_META)
    """Identifier of point, can be shared across stacks."""

    location: Point3D = attr.ib(metadata=SAVE_META)
    """Node position as an (x, y, z) tuple, in pixels."""

    radius: Optional[float] = attr.ib(default=None, metadata=SAVE_META)
    """(optional) radius of the point, in pixels."""

    parentBranch: Optional[Branch] = attr.ib(default=None, repr=False, eq=False, order=False)
    """Branch this point belongs to."""

    annotation: str = attr.ib(default="", eq=False, order=False, metadata=SAVE_META)
    """Text annotation for node."""

    children: List[Branch] = attr.ib(default=attr.Factory(list))
    """Branches coming off the node."""

    manuallyMarked: Optional[bool] = attr.ib(default=None, eq=False, order=False, metadata=SAVE_META)
    """Indicates which points have been marked to be revisited."""

    hilighted: Optional[bool] = attr.ib(default=None, eq=False, order=False, metadata=SAVE_META)
    """ NOTE: Hilighting has been removed, keep here for backwards compatibility."""

    def isRoot(self) -> bool:
        """Whether this point represents the root of the whole tree."""
        return self.parentBranch is None

    def indexInParent(self) -> int:
        """How far along the branch this point sits, 0 = first point after branch point."""
        if self.parentBranch is None:
            return 0 # Root is first
        return self.parentBranch.indexForPoint(self)

    def isLastInBranch(self) -> bool:
        """Return whether this point is the terminal point in the branch."""
        return self.nextPointInBranch(noWrap=True) is None

    def removeChildrenByID(self, branchID: str) -> None:
        """Removes a branch from the list of branches coming off this point."""
        self.children = [b for b in self.children if b.id != branchID]

    def nextPointInBranch(self, delta: int=1, noWrap: bool=False) -> Optional[Point]:
        """Walks a distance along the branch and returns the sibling."""
        if delta == 0:
            return self

        if self.parentBranch is None:
            if delta < 0:
                # Root point has none before
                return None
            else:
                if noWrap:
                    return None
                # Walk down to first branch.
                if len(self.children) > 0 and len(self.children[0].points) > 0:
                    return self.children[0].points[0].nextPointInBranch(delta - 1)
            return None # Not hit - for type hints only

        idx = self.indexInParent()
        nextIdx = idx + delta
        if nextIdx == -1:
            return self.parentBranch.parentPoint
        elif nextIdx >= 0 and nextIdx < len(self.parentBranch.points):
            return self.parentBranch.points[nextIdx]
        if noWrap:
            return None
        # Wrapping allowed, so walk before/after brach.
        if nextIdx < 0:
            if self.parentBranch.parentPoint is not None:
                return self.parentBranch.parentPoint.nextPointInBranch(nextIdx + 1)
            return None
        else:
            assert nextIdx >= len(self.parentBranch.points)
            at = self.parentBranch.points[-1]
            # Overflow, so go down first branch.
            if at is not None and len(at.children) > 0 and len(at.children[0].points) > 0:
                childBranchFirstPoint = at.children[0].points[0]
                return childBranchFirstPoint.nextPointInBranch(nextIdx - len(self.parentBranch.points))
        return None # Not hit - for type hints only

    def flattenSubtreePoints(self) -> List[Point]:
        """Return all points downstream from this point."""
        if self.parentBranch is not None:
            return self.parentBranch.flattenSubtreePoints(startIdx=self.indexInParent())
        # Root point, needs special logic
        points = [self]
        for child in self.children:
            points.extend(child.flattenSubtreePoints())
        return points

    def subtreeContainsID(self, pointID:str) -> bool:
        """Whether the given point ID exists anywhere further down the tree."""
        return pointID in [p.id for p in self.flattenSubtreePoints()]

    def pathFromRoot(self) -> List[Point]:
        """Points in order to get from the root to this point."""
        points = []
        pointAt: Optional[Point] = self
        while pointAt is not None:
            points.append(pointAt)
            pointAt = pointAt.nextPointInBranch(delta=-1)
        return list(reversed(points))

    def radiusFromAncestors(self) -> float:
        """Walks backwards from current point until a radius is found.
            Returns default value (5) if no Radius is found """
        ancestorRadius: Optional[float] = None
        upTreePoint: Optional[Point] = self
        while ancestorRadius == None and upTreePoint is not None:
            upTreePoint = upTreePoint.nextPointInBranch(delta=-1)
            if upTreePoint == None:
                ancestorRadius == 5.0
            if upTreePoint is not None:
                if upTreePoint.radius is not None:
                    ancestorRadius = upTreePoint.radius
            else:
                ancestorRadius = 5.0

        if ancestorRadius is None:
            ancestorRadius = 5.0 # Not hit - for type hints only
        return ancestorRadius

    def longestDistanceToLeaf(self) -> float:
        branchForTree = self.parentBranch
        if branchForTree is None and len(self.children) > 0:
            branchForTree = self.children[0]
        if branchForTree is None:
            # No branches in tree:
            return 0.0

        tree = branchForTree._parentTree

        longestDist = 0.0
        sibling = self.nextPointInBranch(noWrap=True)
        if sibling is not None:
            siblingDist = tree.spatialDist(self, sibling) + sibling.longestDistanceToLeaf()
            longestDist = max(longestDist, siblingDist)
        for child in self.children:
            if len(child.points) > 0:
                cPoint = child.points[0]
                childDist = tree.spatialDist(self, cPoint) + cPoint.longestDistanceToLeaf()
                longestDist = max(longestDist, childDist)
        return longestDist

    def returnWorldRadius(self, fullState) -> float:
        """Returns world radius value of point or Zero if there is no radius.
        :returns: (worldRadius)
        """
        if self.radius is not None:
            xScale, yScale, _ = fullState.projectOptions.pixelSizes
            scaleFactor = math.sqrt(xScale*yScale)
            worldRadius = self.radius * scaleFactor
        else:
            worldRadius = 0

        return worldRadius


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

"""Estimation functions for compute costs of ML models."""

import contextlib
import functools
import re
from typing import Dict, Iterable, List, Optional, Tuple
from absl import flags
from jax._src.lax import lax
import numpy as onp

from aqt.jax import hlo_utils

# pylint: disable=g-direct-tensorflow-import
from tensorflow.compiler.xla.service import hlo_pb2
# pylint: enable=g-direct-tensorflow-import

FLAGS = flags.FLAGS

flags.DEFINE_bool(
    'metadata_enabled',
    default=False,
    help=('Whether to annotate quantization info in HLO metadata.'))

# We intend to use quantization information to estimate the compute cost of
# our ML model. However, currently JAX does not support transferring high-level
# information from the model to its HLO representation, e.g. via metadata.
# Therefore, using the context managers below, this info is appended to the name
# of the op in the original primitives in lax via monkey-patching. After the
# op is instantiated, the original value of the primitive is restored.
# This is not a permanent solution.

# Based on: https://jax.readthedocs.io/en/latest/jax.lax.html
class DotMetadataMonkeyPatch(contextlib.ContextDecorator):
  """Context for passing quantization data to the dot operation."""

  def __init__(self, *, lhs_prec, rhs_prec,
               rhs_is_weight):
    annotation = _quantization_annotation(lhs_prec, rhs_prec, rhs_is_weight)
    self._op_name = 'dot_general_quant' + annotation

  def __enter__(self):
    # pylint: disable=protected-access
    # The following primitive accepts a name argument which is passed into
    # the HLO metadata field. Here, it is the only argument changed from
    # the original lax implementation.
    self._dot_general_p_original = lax.dot_general_p
    lax.dot_general_p = lax.standard_primitive(
        shape_rule=lax._dot_general_shape_rule,
        dtype_rule=lax._dot_general_dtype_rule,
        name=self._op_name,
        translation_rule=lax._dot_general_translation_rule)
    lax.ad.defbilinear(lax.dot_general_p, lax._dot_general_transpose_lhs,
                       lax._dot_general_transpose_rhs)
    lax.batching.primitive_batchers[
        lax.dot_general_p] = lax._dot_general_batch_rule
    lax.masking.masking_rules[lax.dot_general_p] = lax._dot_general_masking_rule
    # pylint: enable=protected-access

  def __exit__(self, *exc):
    # Restore original primitive
    lax.dot_general_p = self._dot_general_p_original

# Based on: https://jax.readthedocs.io/en/latest/jax.lax.html
class ConvMetadataMonkeyPatch(contextlib.ContextDecorator):
  """Context for passing quantization data to the conv operation."""

  def __init__(self, *, weight_prec,
               act_prec):
    annotation = _quantization_annotation(
        act_prec, weight_prec, rhs_is_weight=True)
    self._op_name = 'conv_general_dilated_quant' + annotation

  def __enter__(self):
    # pylint: disable=protected-access
    self._conv_general_dilated_p_original = lax.conv_general_dilated_p
    # The following primitive accepts a name argument which is passed into
    # the HLO metadata field. Here, it is the only argument changed from
    # the original lax implementation.
    lax.conv_general_dilated_p = lax.standard_primitive(
        shape_rule=lax._conv_general_dilated_shape_rule,
        dtype_rule=lax._conv_general_dilated_dtype_rule,
        name=self._op_name,
        translation_rule=functools.partial(
            lax._conv_general_dilated_translation_rule,
            expand_complex_convolutions=False))
    lax.xla.backend_specific_translations['cpu'][
        lax.conv_general_dilated_p] = functools.partial(
            lax._conv_general_dilated_translation_rule,
            expand_complex_convolutions=True)
    lax.xla.backend_specific_translations['gpu'][
        lax.conv_general_dilated_p] = functools.partial(
            lax._conv_general_dilated_translation_rule,
            expand_complex_convolutions=True)
    lax.ad.defbilinear(lax.conv_general_dilated_p,
                       lax._conv_general_dilated_transpose_lhs,
                       lax._conv_general_dilated_transpose_rhs)
    lax.batching.primitive_batchers[
        lax.conv_general_dilated_p] = lax._conv_general_dilated_batch_rule
    lax.masking.masking_rules[
        lax.conv_general_dilated_p] = lax._conv_general_dilated_masking_rule
    # pylint: enable=protected-access

  def __exit__(self, *exc):
    # Restore original primitive
    lax.conv_general_dilated_p = self._conv_general_dilated_p_original

# TODO(abdolrashidi): Add support for QuantOps.FloatQuant for cost estimation.
def _quantization_annotation(lhs_prec, rhs_prec,
                             rhs_is_weight):
  """Returns an annotation to be appended to the name of the quantizable op."""
  bfloat16_prec = 'bf16'

  def _replace_with_bf16_if_prec_is_none(prec):
    return prec if prec is not None else bfloat16_prec

  lhs_prec = _replace_with_bf16_if_prec_is_none(lhs_prec)
  rhs_prec = _replace_with_bf16_if_prec_is_none(rhs_prec)

  quant_annotation = '_lhs{}_rhs{}_lw{}'.format(lhs_prec, rhs_prec,
                                                int(rhs_is_weight))
  return quant_annotation

def _find_lhs_shape(instr,
                    computations):
  """Find the lhs shape of an instruction in HLO computations."""
  for computation in computations:
    for i in computation.instructions:
      # instr.operand_ids[0] contains the input shape.
      if i.id == instr.operand_ids[0]:
        # Input and output batch sizes must match
        assert i.shape.dimensions[0] == instr.shape.dimensions[0]
        return i.shape.dimensions
  return None

def _find_rhs_shape(instr,
                    computations):
  """Find the weight shape of an instruction in HLO computations."""
  for computation in computations:
    for i in computation.instructions:
      # instr.operand_ids[1] contains the weight shape.
      if i.id == instr.operand_ids[1]:
        return i.shape.dimensions
  return None

def _estimate_weights(
    instr,
    computations):
  """Estimate the number of weights in a conv or dot instruction."""
  weight_shape = _find_rhs_shape(instr, computations)
  assert weight_shape is not None
  return onp.prod(weight_shape)

def _estimate_conv_mults(
    instr,
    computations):
  """Estimate the number of multiplications in a convolution instruction."""
  lhs_shape = _find_lhs_shape(instr, computations)
  assert lhs_shape is not None
  input_channels = lhs_shape[-1]
  kernel_size = onp.prod([dim.size for dim in instr.window.dimensions])
  output_conv_dims = [
      instr.shape.dimensions[i]
      for i in instr.convolution_dimension_numbers.output_spatial_dimensions
  ]
  output_image_size = onp.prod(output_conv_dims)
  output_channels = instr.shape.dimensions[
      instr.convolution_dimension_numbers.output_feature_dimension]
  return output_image_size * kernel_size * input_channels * output_channels

def _estimate_dot_mults(
    instr,
    computations):
  """Estimate the number of multiplications in a dot instruction."""
  lhs_shape = _find_lhs_shape(instr, computations)
  assert lhs_shape is not None
  input_channels = lhs_shape[-1]
  output_channels = instr.shape.dimensions[-1]
  return input_channels * output_channels

def _extract_quant_info(
    instr):
  """Extracts lhs and rhs quantization precision from op metadata."""
  if instr.opcode not in _get_supported_ops():
    raise NotImplementedError('Unexpected op detected')
  # If annotated, metadata.op_type would have the following format:
  # '[original op name]_quant_w[weight prec]_a[act prec]'
  if 'quant' not in instr.metadata.op_type:
    raise NotImplementedError('Unable to parse {}'.format(
        instr.metadata.op_type))
  [(lhs_prec_str, rhs_prec_str, rhs_is_weight_str)
  ] = re.findall('_lhs(.*)_rhs(.*)_lw(.*)', instr.metadata.op_type)

  def _extract_prec(prec_str):
    if prec_str.startswith('bf'):
      prec_str = prec_str[2:]
    if prec_str.isnumeric():
      prec = int(prec_str)
    else:
      raise NotImplementedError('Unable to parse {}'.format(
          instr.metadata.op_type))
    return prec

  lhs_prec = _extract_prec(lhs_prec_str)
  rhs_prec = _extract_prec(rhs_prec_str)
  rhs_is_weight = bool(int(rhs_is_weight_str))
  if lhs_prec <= 0 or rhs_prec <= 0:
    raise ValueError('HLO metadata precision annotatation must be a positive '
                     'integer.')
  return lhs_prec, rhs_prec, rhs_is_weight

def _list_supported_ops_from_hlo(
    hlo_proto):
  """Gather and return a list of supported quantizable ops in the HLO."""
  supported_ops = _get_supported_ops()
  target_instructions = []
  computations = hlo_proto.computations
  for computation in computations:
    for instr in computation.instructions:
      if instr.opcode in supported_ops:
        target_instructions.append(instr)
  return target_instructions

def _get_supported_ops():
  """Get the supported ops for compute and memory cost estimation."""
  # Output dictionary key is HLO instruction opcode.
  # 'estimate_instr_mult': function used to estimate the number of matrix
  # multiplications.
  # 'estimate_instr_weights': function used to estimate the number of weights.
  return {
      'convolution': {
          'estimate_instr_mult': _estimate_conv_mults,
          'estimate_instr_weights': _estimate_weights,
      },
      'dot': {
          'estimate_instr_mult': _estimate_dot_mults,
          'estimate_instr_weights': _estimate_weights,
      },
  }

def estimate_compute_cost(
    hlo_proto):
  """Estimates the compute cost for the input HLO proto.

  Args:
    hlo_proto: the model's HLO representation (e.g. derived from
      hlo_utils.load_hlo_proto()) It contains all the instructions used in the
      model.

  Returns:
    a dictionary with two key-value pairs:
      'compute_cost': The sum of all compute costs for each layer
      'compute_cost_ratio_to_bfloat16': The ratio of the estimated compute cost
      to the cost in the case of no quantization (bfloat16)

  """
  # To estimate the overall compute cost of ops, we should multiply the
  # number of multiplications by bits_weights * bits_acts in each layer.

  # Gather the supported quantizable ops and their quantization parameters
  target_instructions = _list_supported_ops_from_hlo(hlo_proto)

  # Begin compute cost calculation
  compute_cost_quadratic = 0
  compute_cost_linear = 0
  bfloat16_cost_quadratic = 0  # computed as reference for comparison purposes
  bfloat16_cost_linear = 0

  supported_ops = _get_supported_ops()
  for instr in target_instructions:
    # Estimate multiplications in the op
    opcode = instr.opcode
    multiplication_count = supported_ops[opcode]['estimate_instr_mult'](
        instr, hlo_proto.computations)

    # For the model's compute cost, we use the number of multiplications used
    # in the quantizable layers used in the model, which dominate the cost
    # compared to that of other operations used in them, such as additions.
    # Also, the cost of each multiplication is proportional to the number of
    # bits in each operand, which are subject to change due to quantization.
    # Therefore, we use the weight and activation precisions in our estimation
    # as well.

    bfloat16_lhs_prec = 16
    bfloat16_rhs_prec = 16

    lhs_prec, rhs_prec, _ = _extract_quant_info(instr)
    if lhs_prec > 16:
      raise ValueError(f'Unexpected lhs precision {lhs_prec}.')
    if rhs_prec > 16:
      raise ValueError(f'Unexpected rhs precision {rhs_prec}.')

    bfloat16_cost_quadratic += multiplication_count * (
        bfloat16_lhs_prec * bfloat16_rhs_prec)
    compute_cost_quadratic += multiplication_count * lhs_prec * rhs_prec

    # Nvidia A100 only supports 16x16, 8x8 and 4x4
    # Cost of 8x8 is half of 16x16
    # Cost of 4x4 is 1/4 of 16x16
    bfloat16_cost_linear += multiplication_count * max(
        bfloat16_lhs_prec, bfloat16_rhs_prec)

    prec = max(lhs_prec, rhs_prec)
    if prec <= 4:
      prec = 4
    elif prec <= 8:
      prec = 8
    elif prec <= 16:
      prec = 16

    compute_cost_linear += multiplication_count * prec
  # Return the results
  cost_ratio_to_bfloat16_quadratic = compute_cost_quadratic / bfloat16_cost_quadratic
  cost_ratio_to_bfloat16_linear = compute_cost_linear / bfloat16_cost_linear

  result_dict = {
      # quadratic compute cost. We didn't rename the keys to
      # compute_cost_quadratic to presev backwards compatibility.
      'compute_cost':
          float(compute_cost_quadratic),
      'compute_cost_ratio_to_bfloat16':
          float(cost_ratio_to_bfloat16_quadratic),
      # Linear compute cost
      'compute_cost_linear':
          float(compute_cost_linear),
      'compute_cost_ratio_to_bfloat16_linear':
          float(cost_ratio_to_bfloat16_linear),
  }

  return result_dict

def estimate_memory_cost(hlo_proto):
  """Estimates the memory cost for the input HLO proto.

  Args:
    hlo_proto: the model's HLO representation (e.g. derived from
      hlo_utils.load_hlo_proto()) It contains all the instructions used in the
      model.

  Returns:
    a dictionary with two key-value pairs:
      'memory_cost': The sum of all memory costs for each layer
      'memory_cost_ratio_to_bfloat16': The ratio of the estimated memory cost to
         the cost in the case of no quantization (bfloat16)

  """
  # For the memory cost, we estimate the number of bits used for the weights,
  # which would also require the quantization precision for those weights.

  # Gather the supported quantizable ops and their quantization parameters
  target_instructions = _list_supported_ops_from_hlo(hlo_proto)

  # Begin memory cost calculation
  memory_cost = 0
  bfloat16_cost = 0  # computed as reference for comparison purposes

  supported_ops = _get_supported_ops()
  for instr in target_instructions:
    # Estimate number of weights in the op
    opcode = instr.opcode
    weight_count = supported_ops[opcode]['estimate_instr_weights'](
        instr, hlo_proto.computations)

    # Multiply number of weights by the number of bits used for each.
    bfloat16_weight_prec = 16

    _, rhs_prec, rhs_is_weight = _extract_quant_info(instr)

    if rhs_prec > 16:
      raise ValueError(f'Unexpected rhs precision {rhs_prec}.')

    if rhs_is_weight:
      memory_cost += weight_count * rhs_prec
      bfloat16_cost += weight_count * bfloat16_weight_prec

  # For dynamic matmuls, memory_cost is zero. The following check is to avoid
  # division by zero. We set the ratio for 0/0 to 1.
  if bfloat16_cost == 0 and memory_cost == 0:
    cost_ratio_to_bfloat16 = 1
  else:
    cost_ratio_to_bfloat16 = memory_cost / bfloat16_cost

  result_dict = {
      'memory_cost': float(memory_cost),
      'memory_cost_ratio_to_bfloat16': float(cost_ratio_to_bfloat16)
  }
  return result_dict

def estimate_costs_of_dot_and_conv_ops_from_jax_fn(
    fn, *fn_args, **fn_kwargs):
  """Wrapper function around estimate_compute_and_memory_cost().

  Will generate hlo proto from jax function, and call
  estimate_compute_and_memory_cost() on it.

  Args:
    fn: the function for which the HLO is to be produced.
    *fn_args: the function's args.
    **fn_kwargs: the function's kwargs.

  Returns:
    A dictionary with compute cost and memory cost data.
    See estimate_compute_and_memory_cost() docstring for details.
  """
  FLAGS.metadata_enabled = True
  hlo_module_proto = hlo_utils.load_hlo_proto_from_jax_fn(
      fn, *fn_args, **fn_kwargs)
  cost_dict = estimate_compute_cost(hlo_module_proto)
  memory_cost_dict = estimate_memory_cost(hlo_module_proto)
  cost_dict.update(memory_cost_dict)
  FLAGS.metadata_enabled = False
  return cost_dict


import ctypes

import OpenGL.GL as GL
import numpy as np

from PIL import Image

skyBoxVertices = np.array([-1.0, 1.0, -1.0,
                           -1.0, -1.0, -1.0,
                           1.0, -1.0, -1.0,
                           1.0, -1.0, -1.0,
                           1.0, 1.0, -1.0,
                           -1.0, 1.0, -1.0,

                           -1.0, -1.0, 1.0,
                           -1.0, -1.0, -1.0,
                           -1.0, 1.0, -1.0,
                           -1.0, 1.0, -1.0,
                           -1.0, 1.0, 1.0,
                           -1.0, -1.0, 1.0,

                           1.0, -1.0, -1.0,
                           1.0, -1.0, 1.0,
                           1.0, 1.0, 1.0,
                           1.0, 1.0, 1.0,
                           1.0, 1.0, -1.0,
                           1.0, -1.0, -1.0,

                           -1.0, -1.0, 1.0,
                           -1.0, 1.0, 1.0,
                           1.0, 1.0, 1.0,
                           1.0, 1.0, 1.0,
                           1.0, -1.0, 1.0,
                           -1.0, -1.0, 1.0,

                           -1.0, 1.0, -1.0,
                           1.0, 1.0, -1.0,
                           1.0, 1.0, 1.0,
                           1.0, 1.0, 1.0,
                           -1.0, 1.0, 1.0,
                           -1.0, 1.0, -1.0,

                           -1.0, -1.0, -1.0,
                           -1.0, -1.0, 1.0,
                           1.0, -1.0, -1.0,
                           1.0, -1.0, -1.0,
                           -1.0, -1.0, 1.0,
                           1.0, -1.0, 1.0], dtype=np.float32)

class SkyBox:
    def __init__(self, sky_box_dir: str, ext: str):
        """
        :param sky_box_dir: must contain files ("right", "left", "top", "bottom", "front", "back").ext
        :param ext: one of png or jpg
        """
        self.VAO = GL.glGenVertexArrays(1)
        GL.glBindVertexArray(self.VAO)

        self.VBO = GL.glGenBuffers(1)
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.VBO)
        GL.glBufferData(GL.GL_ARRAY_BUFFER, skyBoxVertices.nbytes, skyBoxVertices, GL.GL_STATIC_DRAW)

        GL.glEnableVertexAttribArray(0)
        GL.glVertexAttribPointer(0, 3, GL.GL_FLOAT, GL.GL_FALSE, 3 * 4, ctypes.c_void_p(0))

        self.cubeMap = GL.glGenTextures(1)
        GL.glBindTexture(GL.GL_TEXTURE_CUBE_MAP, self.cubeMap)

        skyBoxImages = ("right", "left", "top", "bottom", "front", "back")

        for i in range(6):
            img = Image.open(sky_box_dir + f"/{skyBoxImages[i]}.{ext}")
            GL.glTexImage2D(GL.GL_TEXTURE_CUBE_MAP_POSITIVE_X + i, 0, GL.GL_RGB, img.width, img.height,
                            0, GL.GL_RGB, GL.GL_UNSIGNED_BYTE,
                            img.convert("RGB").tobytes())

        GL.glTexParameteri(GL.GL_TEXTURE_CUBE_MAP, GL.GL_TEXTURE_MAG_FILTER, GL.GL_LINEAR)
        GL.glTexParameteri(GL.GL_TEXTURE_CUBE_MAP, GL.GL_TEXTURE_MIN_FILTER, GL.GL_LINEAR)
        GL.glTexParameteri(GL.GL_TEXTURE_CUBE_MAP, GL.GL_TEXTURE_WRAP_S, GL.GL_CLAMP_TO_EDGE)
        GL.glTexParameteri(GL.GL_TEXTURE_CUBE_MAP, GL.GL_TEXTURE_WRAP_T, GL.GL_CLAMP_TO_EDGE)
        GL.glTexParameteri(GL.GL_TEXTURE_CUBE_MAP, GL.GL_TEXTURE_WRAP_R, GL.GL_CLAMP_TO_EDGE)

    def draw(self):
        GL.glBindVertexArray(self.VAO)
        # bind texture to unit 0 in order to avoid setting uniform
        GL.glActiveTexture(GL.GL_TEXTURE0)
        GL.glBindTexture(GL.GL_TEXTURE_CUBE_MAP, self.cubeMap)

        GL.glDepthMask(GL.GL_FALSE)
        GL.glDrawArrays(GL.GL_TRIANGLES, 0, 36)
        GL.glDepthMask(GL.GL_TRUE)


import torch
import numpy as np
import logging, yaml, os, sys, argparse, time, math
from tqdm import tqdm
from collections import defaultdict
from Logger import Logger
import matplotlib
matplotlib.use('agg')
matplotlib.rcParams['agg.path.chunksize'] = 10000
import matplotlib.pyplot as plt
from scipy.io import wavfile
from random import choice

from Modules import PVCGAN, MultiResolutionSTFTLoss
from Datasets import Train_Dataset, Dev_Dataset, Accumulation_Dataset, Inference_Dataset, Collater, Accumulation_Collater, Inference_Collater
from Radam import RAdam
from Noam_Scheduler import Modified_Noam_Scheduler

from Pattern_Generator import Pattern_Generate
from yin import pitch_calc

with open('Hyper_Parameters.yaml') as f:
    hp_Dict = yaml.load(f, Loader=yaml.Loader)

if not hp_Dict['Device'] is None:
    os.environ['CUDA_VISIBLE_DEVICES']= hp_Dict['Device']

if not torch.cuda.is_available():
    device = torch.device('cpu')
else:
    device = torch.device('cuda:0')
    torch.backends.cudnn.benchmark = True
    torch.cuda.set_device(0)

logging.basicConfig(
    level=logging.INFO, # level=logging.DEBUG,
    stream=sys.stdout,
    format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
    )

if hp_Dict['Use_Mixed_Precision']:
    try:
        from apex import amp
    except:
        logging.info('There is no apex modules in the environment. Mixed precision does not work.')
        hp_Dict['Use_Mixed_Precision'] = False

class Trainer:
    def __init__(self, steps= 0):
        self.steps = steps
        self.epochs = 0

        self.Datset_Generate()
        self.Model_Generate()

        self.scalar_Dict = {
            'Train': defaultdict(float),
            'Evaluation': defaultdict(float),
            }

        self.writer_Dict = {
            'Train': Logger(os.path.join(hp_Dict['Log_Path'], 'Train')),
            'Evaluation': Logger(os.path.join(hp_Dict['Log_Path'], 'Evaluation')),
            }

        self.Load_Checkpoint()

    def Datset_Generate(self):
        train_Dataset = Train_Dataset()
        accumulation_Dataset = Accumulation_Dataset()
        dev_Dataset = Dev_Dataset()
        inference_Dataset = Inference_Dataset()
        logging.info('The number of base train files = {}.'.format(len(train_Dataset) // hp_Dict['Train']['Train_Pattern']['Accumulated_Dataset_Epoch']))
        logging.info('The number of development patterns = {}.'.format(len(dev_Dataset)))
        logging.info('The number of inference patterns = {}.'.format(len(inference_Dataset)))

        collater = Collater()
        accumulation_Collater = Accumulation_Collater()
        inference_Collater = Inference_Collater()

        self.dataLoader_Dict = {}
        self.dataLoader_Dict['Train'] = torch.utils.data.DataLoader(
            dataset= train_Dataset,
            shuffle= True,
            collate_fn= collater,
            batch_size= hp_Dict['Train']['Batch_Size'],
            num_workers= hp_Dict['Train']['Num_Workers'],
            pin_memory= True
            )
        self.dataLoader_Dict['Accumulation'] = torch.utils.data.DataLoader(
            dataset= accumulation_Dataset,
            shuffle= False,
            collate_fn= accumulation_Collater,
            batch_size= hp_Dict['Train']['Batch_Size'],
            num_workers= hp_Dict['Train']['Num_Workers'],
            pin_memory= True
            )
        self.dataLoader_Dict['Dev'] = torch.utils.data.DataLoader(
            dataset= dev_Dataset,
            shuffle= True,  # to write tensorboard.
            collate_fn= collater,
            batch_size= hp_Dict['Train']['Batch_Size'],
            num_workers= hp_Dict['Train']['Num_Workers'],
            pin_memory= True
            )
        self.dataLoader_Dict['Inference'] = torch.utils.data.DataLoader(
            dataset= inference_Dataset,
            shuffle= False,  # to write tensorboard.
            collate_fn= inference_Collater,
            batch_size= hp_Dict['Inference_Batch_Size'] or hp_Dict['Train']['Batch_Size'],
            num_workers= hp_Dict['Train']['Num_Workers'],
            pin_memory= True
            )

    def Model_Generate(self):
        self.model = PVCGAN().to(device)
        self.criterion_Dict = {
            'STFT': MultiResolutionSTFTLoss(
                fft_sizes= hp_Dict['STFT_Loss_Resolution']['FFT_Sizes'],
                shift_lengths= hp_Dict['STFT_Loss_Resolution']['Shfit_Lengths'],
                win_lengths= hp_Dict['STFT_Loss_Resolution']['Win_Lengths'],
                ).to(device),
            'MSE': torch.nn.MSELoss().to(device),
            'CE': torch.nn.CrossEntropyLoss().to(device),
            'MAE': torch.nn.L1Loss().to(device)
            }

        self.optimizer = RAdam(
            params= self.model.parameters(),
            lr= hp_Dict['Train']['Learning_Rate']['Initial'],
            betas=(hp_Dict['Train']['ADAM']['Beta1'], hp_Dict['Train']['ADAM']['Beta2']),
            eps= hp_Dict['Train']['ADAM']['Epsilon'],
            weight_decay= hp_Dict['Train']['Weight_Decay']
            )
        self.scheduler = Modified_Noam_Scheduler(
            optimizer= self.optimizer,
            base= hp_Dict['Train']['Learning_Rate']['Base']
            )

        if hp_Dict['Use_Mixed_Precision']:
            self.model, self.optimizer = amp.initialize(
                models=self.model,
                optimizers=self.optimizer
                )

        logging.info(self.model)

    def Train_Step(self, audios, mels, pitches, audio_Singers, mel_Singers, noises):
        loss_Dict = {}

        audios = audios.to(device, non_blocking=True)
        mels = mels.to(device, non_blocking=True)
        pitches = pitches.to(device, non_blocking=True)
        audio_Singers = audio_Singers.to(device, non_blocking=True)
        mel_Singers = mel_Singers.to(device, non_blocking=True)
        noises = noises.to(device, non_blocking=True)

        # Generator
        fakes, predicted_Singers, predicted_Pitches, fakes_Discriminations, reals_Discriminations = self.model(
            mels= mels,
            pitches= pitches,
            singers= audio_Singers,
            noises= noises,
            discrimination= self.steps >= hp_Dict['Train']['Discriminator_Delay'],
            reals= audios
            )

        loss_Dict['Generator/Spectral_Convergence'], loss_Dict['Generator/Magnitude'] = self.criterion_Dict['STFT'](fakes, audios)
        loss_Dict['Generator'] = loss_Dict['Generator/Spectral_Convergence'] + loss_Dict['Generator/Magnitude']

        loss_Dict['Confuser/Singer'] = self.criterion_Dict['CE'](predicted_Singers, mel_Singers)
        loss_Dict['Confuser/Pitch'] = self.criterion_Dict['MAE'](predicted_Pitches, pitches)
        loss_Dict['Confuser'] =  loss_Dict['Confuser/Singer'] + loss_Dict['Confuser/Pitch']
        loss = loss_Dict['Generator'] + loss_Dict['Confuser']

        if self.steps >= hp_Dict['Train']['Discriminator_Delay']:
            loss_Dict['Discriminator/Fake'] = self.criterion_Dict['MSE'](fakes_Discriminations, torch.zeros_like(fakes_Discriminations))   # Discriminator thinks that 0 is correct.
            loss_Dict['Discriminator/Real'] = self.criterion_Dict['MSE'](reals_Discriminations, torch.ones_like(reals_Discriminations))  # Discriminator thinks that 1 is correct.

            loss_Dict['Discriminator'] = loss_Dict['Discriminator/Fake'] + loss_Dict['Discriminator/Real']
            loss += loss_Dict['Discriminator']

        self.optimizer.zero_grad()

        if hp_Dict['Use_Mixed_Precision']:
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                scaled_loss.backward()
            torch.nn.utils.clip_grad_norm_(
                parameters= amp.master_params(self.optimizer),
                max_norm= hp_Dict['Train']['Gradient_Norm']
                )
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                parameters= self.model.parameters(),
                max_norm= hp_Dict['Train']['Gradient_Norm']
                )

        self.optimizer.step()
        self.scheduler.step()

        self.steps += 1
        self.tqdm.update(1)

        for tag, loss in loss_Dict.items():
            self.scalar_Dict['Train']['Loss/{}'.format(tag)] += loss

    def Train_Epoch(self):
        if any([
            hp_Dict['Train']['Train_Pattern']['Mixup']['Use'] and self.steps >= hp_Dict['Train']['Train_Pattern']['Mixup']['Apply_Delay'],
            hp_Dict['Train']['Train_Pattern']['Back_Translate']['Use'] and self.steps >= hp_Dict['Train']['Train_Pattern']['Back_Translate']['Apply_Delay']
            ]):
            self.Data_Accumulation()

        for audios, mels, pitches, audio_Singers, mel_Singers, noises in self.dataLoader_Dict['Train']:
            self.Train_Step(audios, mels, pitches, audio_Singers, mel_Singers, noises)

            if self.steps % hp_Dict['Train']['Checkpoint_Save_Interval'] == 0:
                self.Save_Checkpoint()

            if self.steps % hp_Dict['Train']['Logging_Interval'] == 0:
                self.scalar_Dict['Train'] = {
                    tag: loss / hp_Dict['Train']['Logging_Interval']
                    for tag, loss in self.scalar_Dict['Train'].items()
                        }
                self.scalar_Dict['Train']['Learning_Rate'] = self.scheduler.get_last_lr()
                self.writer_Dict['Train'].add_scalar_dict(self.scalar_Dict['Train'], self.steps)
                self.scalar_Dict['Train'] = defaultdict(float)

            if self.steps % hp_Dict['Train']['Evaluation_Interval'] == 0:
                self.Evaluation_Epoch()

            if self.steps % hp_Dict['Train']['Inference_Interval'] == 0:
                self.Inference_Epoch()

            if self.steps >= hp_Dict['Train']['Max_Step']:
                return

        self.epochs += hp_Dict['Train']['Train_Pattern']['Accumulated_Dataset_Epoch']

    @torch.no_grad()
    def Evaluation_Step(self, audios, mels, pitches, audio_Singers, mel_Singers, noises):
        loss_Dict = {}

        audios = audios.to(device, non_blocking=True)
        mels = mels.to(device, non_blocking=True)
        pitches = pitches.to(device, non_blocking=True)
        audio_Singers = audio_Singers.to(device, non_blocking=True)
        mel_Singers = mel_Singers.to(device, non_blocking=True)
        noises = noises.to(device, non_blocking=True)

        # Generator
        fakes, predicted_Singers, predicted_Pitches, fakes_Discriminations, reals_Discriminations = self.model(
            mels= mels,
            pitches= pitches,
            singers= audio_Singers,
            noises= noises,
            discrimination= self.steps >= hp_Dict['Train']['Discriminator_Delay'],
            reals= audios
            )

        loss_Dict['Generator/Spectral_Convergence'], loss_Dict['Generator/Magnitude'] = self.criterion_Dict['STFT'](fakes, audios)
        loss_Dict['Generator'] = loss_Dict['Generator/Spectral_Convergence'] + loss_Dict['Generator/Magnitude']

        loss_Dict['Confuser/Singer'] = self.criterion_Dict['CE'](predicted_Singers, mel_Singers)
        loss_Dict['Confuser/Pitch'] = self.criterion_Dict['MAE'](predicted_Pitches, pitches)
        loss_Dict['Confuser'] =  loss_Dict['Confuser/Singer'] + loss_Dict['Confuser/Pitch']
        loss = loss_Dict['Generator'] + loss_Dict['Confuser']

        if self.steps >= hp_Dict['Train']['Discriminator_Delay']:
            loss_Dict['Discriminator/Fake'] = self.criterion_Dict['MSE'](fakes_Discriminations, torch.zeros_like(fakes_Discriminations))   # Discriminator thinks that 0 is correct.
            loss_Dict['Discriminator/Real'] = self.criterion_Dict['MSE'](reals_Discriminations, torch.ones_like(reals_Discriminations))  # Discriminator thinks that 1 is correct.

            loss_Dict['Discriminator'] = loss_Dict['Discriminator/Fake'] + loss_Dict['Discriminator/Real']
            loss += loss_Dict['Discriminator']

        for tag, loss in loss_Dict.items():
            self.scalar_Dict['Evaluation']['Loss/{}'.format(tag)] += loss

        return fakes, predicted_Singers, predicted_Pitches

    def Evaluation_Epoch(self):
        logging.info('(Steps: {}) Start evaluation.'.format(self.steps))

        self.model.eval()

        for step, (audios, mels, pitches, audio_Singers, mel_Singers, noises) in tqdm(
            enumerate(self.dataLoader_Dict['Dev'], 1),
            desc='[Evaluation]',
            total= math.ceil(len(self.dataLoader_Dict['Dev'].dataset) / hp_Dict['Train']['Batch_Size'])
            ):
            fakes, predicted_Singers, predicted_Pitches = self.Evaluation_Step(audios, mels, pitches, audio_Singers, mel_Singers, noises)

        self.scalar_Dict['Evaluation'] = {
            tag: loss / step
            for tag, loss in self.scalar_Dict['Evaluation'].items()
            }
        self.writer_Dict['Evaluation'].add_scalar_dict(self.scalar_Dict['Evaluation'], self.steps)
        self.writer_Dict['Evaluation'].add_histogram_model(self.model, self.steps, delete_keywords=['layer_Dict', '1', 'layer'])
        self.scalar_Dict['Evaluation'] = defaultdict(float)

        self.writer_Dict['Evaluation'].add_image_dict({
            'Mel': (mels[-1].cpu().numpy(), None),
            'Audio/Original': (audios[-1].cpu().numpy(), None),
            'Audio/Predicted': (fakes[-1].cpu().numpy(), None),
            'Pitch/Original': (pitches[-1].cpu().numpy(), None),
            'Pitch/Predicted': (predicted_Pitches[-1].cpu().numpy(), None),
            'Singer/Original': (torch.nn.functional.one_hot(mel_Singers, hp_Dict['Num_Singers']).cpu().numpy(), None),
            'Singer/Predicted': (torch.softmax(predicted_Singers, dim= -1).cpu().numpy(), None),
            }, self.steps)

        self.model.train()

    @torch.no_grad()
    def Inference_Step(self, audios, mels, pitches, singers, noises, source_Labels, singer_Labels, start_Index= 0, tag_step= False, tag_Index= False):
        audios = audios.to(device, non_blocking=True)
        mels = mels.to(device, non_blocking=True)
        pitches = pitches.to(device, non_blocking=True)
        singers = singers.to(device, non_blocking=True)
        noises = noises.to(device, non_blocking=True)

        fakes, *_ = self.model(
            mels= mels,
            pitches= pitches,
            singers= singers,
            noises= noises
            )

        files = []
        for index, (source_Label, singer_Label) in enumerate(zip(source_Labels, singer_Labels)):
            tags = []
            if tag_step: tags.append('Step-{}'.format(self.steps))
            tags.append('{}_to_{}'.format(source_Label, singer_Label))
            if tag_Index: tags.append('IDX_{}'.format(index + start_Index))
            files.append('.'.join(tags))

        os.makedirs(os.path.join(hp_Dict['Inference_Path'], 'Step-{}'.format(self.steps), 'PNG').replace('\\', '/'), exist_ok= True)
        os.makedirs(os.path.join(hp_Dict['Inference_Path'], 'Step-{}'.format(self.steps), 'WAV').replace("\\", "/"), exist_ok= True)
        for index, (real, fake, source_Label, singer_Label, file) in enumerate(zip(
            audios.cpu().numpy(),
            fakes.cpu().numpy(),
            source_Labels,
            singer_Labels,
            files
            )):
            new_Figure = plt.figure(figsize=(80, 10 * 2), dpi=100)
            plt.subplot(211)
            plt.plot(real)
            plt.title('Original wav    Index: {}    {} -> {}'.format(index + start_Index, source_Label, singer_Label))
            plt.subplot(212)
            plt.plot(fake)
            plt.title('Fake wav    Index: {}    {} -> {}'.format(index + start_Index, source_Label, singer_Label))
            plt.tight_layout()
            plt.savefig(os.path.join(hp_Dict['Inference_Path'], 'Step-{}'.format(self.steps), 'PNG', '{}.png'.format(file)).replace('\\', '/'))
            plt.close(new_Figure)

            wavfile.write(
                filename= os.path.join(hp_Dict['Inference_Path'], 'Step-{}'.format(self.steps), 'WAV', '{}.wav'.format(file)).replace('\\', '/'),
                data= (np.clip(fake, -1.0 + 1e-7, 1.0 - 1e-7) * 32767.5).astype(np.int16),
                rate= hp_Dict['Sound']['Sample_Rate']
                )

    def Inference_Epoch(self):
        logging.info('(Steps: {}) Start inference.'.format(self.steps))

        self.model.eval()

        for step, (audios, mels, pitches, singers, noises, source_Labels, singer_Labels) in tqdm(
            enumerate(self.dataLoader_Dict['Inference'], 1),
            desc='[Inference]',
            total= math.ceil(len(self.dataLoader_Dict['Inference'].dataset) / (hp_Dict['Inference_Batch_Size'] or hp_Dict['Train']['Batch_Size']))
            ):
            self.Inference_Step(audios, mels, pitches, singers, noises, source_Labels, singer_Labels, start_Index= (step - 1) * hp_Dict['Train']['Batch_Size'])

        self.model.train()

    @torch.no_grad()
    def Back_Translate_Step(self, mels, pitches, singers, noises):
        mels = mels.to(device, non_blocking=True)
        pitches = pitches.to(device, non_blocking=True)
        singers = singers.to(device, non_blocking=True)
        noises = noises.to(device, non_blocking=True)

        fakes, *_ = self.model(
            mels= mels,
            pitches= pitches,
            singers= singers,
            noises= noises
            )

        return fakes.cpu().numpy()

    def Data_Accumulation(self):
        def Mixup(audio, pitch):
            max_Offset = pitch.shape[0] - hp_Dict['Train']['Wav_Length'] // hp_Dict['Sound']['Frame_Shift'] * 2
            offset1 = np.random.randint(0, max_Offset)
            offset2 = np.random.randint(0, max_Offset)
            beta = np.random.uniform(
                low= hp_Dict['Train']['Train_Pattern']['Mixup']['Min_Beta'],
                high= hp_Dict['Train']['Train_Pattern']['Mixup']['Max_Beta'],
                )

            new_Audio = \
                beta * audio[offset1 * hp_Dict['Sound']['Frame_Shift']:offset1 * hp_Dict['Sound']['Frame_Shift'] + hp_Dict['Train']['Wav_Length'] * 2] + \
                (1 - beta) * audio[offset2* hp_Dict['Sound']['Frame_Shift']:offset2 * hp_Dict['Sound']['Frame_Shift'] + hp_Dict['Train']['Wav_Length'] * 2]

            new_Pitch = \
                beta * pitch[offset1:offset1 + hp_Dict['Train']['Wav_Length'] // hp_Dict['Sound']['Frame_Shift'] * 2] + \
                (1 - beta) * pitch[offset2:offset2 + hp_Dict['Train']['Wav_Length'] // hp_Dict['Sound']['Frame_Shift'] * 2]

            _, new_Mel, _, _ = Pattern_Generate(audio= new_Audio)

            return new_Audio, new_Mel, new_Pitch

        def Back_Translate(mels, pitches, singers, noises):
            fakes = self.Back_Translate_Step(
                mels= mels,
                pitches= pitches,
                singers= singers,
                noises= noises
                )

            new_Mels = [
                Pattern_Generate(audio= fake)[1]
                for fake in fakes
                ]

            return new_Mels

        print()
        mixup_List = []
        back_Translate_List = []
        for total_Audios, total_Pitches, audios, mels, pitches, singers, noises in tqdm(
            self.dataLoader_Dict['Accumulation'],
            desc='[Accumulation]',
            total= math.ceil(len(self.dataLoader_Dict['Accumulation'].dataset) / hp_Dict['Train']['Batch_Size'])
            ):
            #Mixup
            if hp_Dict['Train']['Train_Pattern']['Mixup']['Use'] and self.steps >= hp_Dict['Train']['Train_Pattern']['Mixup']['Apply_Delay']:
                for audio, pitch, singer in zip(total_Audios, total_Pitches, singers.numpy()):
                    mixup_Audio, mixup_Mel, mixup_Pitch = Mixup(audio, pitch)
                    mixup_List.append((
                        mixup_Audio,
                        mixup_Mel,
                        mixup_Pitch,
                        singer,
                        singer
                        ))

            #Backtranslate
            if hp_Dict['Train']['Train_Pattern']['Back_Translate']['Use'] and self.steps >= hp_Dict['Train']['Train_Pattern']['Back_Translate']['Apply_Delay']:
                mel_Singers = torch.LongTensor(np.stack([
                    choice([x for x in range(hp_Dict['Num_Singers']) if x != singer])
                    for singer in singers
                    ], axis= 0))
                back_Translate_Mels = Back_Translate(mels, pitches, mel_Singers, noises)
                for audio, back_Translate_Mel, pitch, audio_Singer, mel_Singer in zip(
                    audios.numpy(), back_Translate_Mels, pitches.numpy(), singers.numpy(), mel_Singers.numpy()):
                    back_Translate_List.append((
                        audio,
                        back_Translate_Mel,
                        pitch,
                        audio_Singer,
                        mel_Singer
                        ))

        self.dataLoader_Dict['Train'].dataset.Accumulation_Renew(
            mixup_Pattern_List= mixup_List,
            back_Translate_Pattern_List= back_Translate_List
            )

    def Load_Checkpoint(self):
        if self.steps == 0:
            paths = [
                os.path.join(root, file).replace('\\', '/')
                for root, _, files in os.walk(hp_Dict['Checkpoint_Path'])
                for file in files
                if os.path.splitext(file)[1] == '.pt'
                ]
            if len(paths) > 0:
                path = max(paths, key = os.path.getctime)
            else:
                return  # Initial training
        else:
            path = os.path.join(hp_Dict['Checkpoint_Path'], 'S_{}.pt'.format(self.steps).replace('\\', '/'))

        state_Dict = torch.load(path, map_location= 'cpu')
        self.model.load_state_dict(state_Dict['Model'])
        self.optimizer.load_state_dict(state_Dict['Optimizer'])
        self.scheduler.load_state_dict(state_Dict['Scheduler'])
        self.steps = state_Dict['Steps']
        self.epochs = state_Dict['Epochs']

        if hp_Dict['Use_Mixed_Precision']:
            if not 'AMP' in state_Dict.keys():
                logging.info('No AMP state dict is in the checkpoint. Model regards this checkpoint is trained without mixed precision.')
            else:
                amp.load_state_dict(state_Dict['AMP'])

        logging.info('Checkpoint loaded at {} steps.'.format(self.steps))

    def Save_Checkpoint(self):
        os.makedirs(hp_Dict['Checkpoint_Path'], exist_ok= True)

        state_Dict = {
            'Model' : self.model.state_dict(),
            'Optimizer': self.optimizer.state_dict(),
            'Scheduler': self.scheduler.state_dict(),
            'Steps': self.steps,
            'Epochs': self.epochs,
            }
        if hp_Dict['Use_Mixed_Precision']:
            state_Dict['AMP'] = amp.state_dict()

        torch.save(
            state_Dict,
            os.path.join(hp_Dict['Checkpoint_Path'], 'S_{}.pt'.format(self.steps).replace('\\', '/'))
            )

        logging.info('Checkpoint saved at {} steps.'.format(self.steps))

    def Train(self):
        if not os.path.exists(os.path.join(hp_Dict['Checkpoint_Path'], 'Hyper_Parameters.yaml')):
            os.makedirs(hp_Dict['Checkpoint_Path'], exist_ok= True)
            with open(os.path.join(hp_Dict['Checkpoint_Path'], 'Hyper_Parameters.yaml').replace("\\", "/"), "w") as f:
                yaml.dump(hp_Dict, f)

        if self.steps == 0:
            self.Evaluation_Epoch()

        if hp_Dict['Train']['Initial_Inference']:
            self.Inference_Epoch()

        self.tqdm = tqdm(
            initial= self.steps,
            total= hp_Dict['Train']['Max_Step'],
            desc='[Training]'
            )

        while self.steps < hp_Dict['Train']['Max_Step']:
            try:
                self.Train_Epoch()
            except KeyboardInterrupt:
                self.Save_Checkpoint()
                exit(1)

        self.tqdm.close()
        logging.info('Finished training.')

if __name__ == '__main__':
    argParser = argparse.ArgumentParser()
    argParser.add_argument('-s', '--steps', default= 0, type= int)
    args = argParser.parse_args()

    new_Trainer = Trainer(steps= args.steps)
    new_Trainer.Train()


import pandas as pd
import numpy as np
import nltk
from nltk.tokenize import word_tokenize
from spacy.lang.en.stop_words import STOP_WORDS
import re
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfVectorizer
from nltk.tokenize import RegexpTokenizer
from gensim.models import doc2vec
from collections import namedtuple
# hide settingwithcopywarning
pd.options.mode.chained_assignment = None

def rank(jobID, topX, root_file_path, all_resumes):

    # read structure + one hot encoded dfs
    job_dummies_ideal = pd.read_csv(root_file_path + "Resume-Parser-JOBS/data/job_description_one_hot_ideal_FULL.csv")
    resume_dummies = pd.read_csv(root_file_path + "Resume-Parser-master-new/data/output/resume_summary_one_hot.csv")

    # cols = ['ReqID', 'text', 'Language', 'Work', 'Summaries', 'Skill', 'Member',
    #    'Writing', 'Researching', 'Honor', 'Activity', 'Education',
    #    'Extracurriculars', 'email', 'phone', 'GPA', 'years_experience',
    #    'mos_experience', 'certifications', 'honor_societies', 'latin_honors',
    #    'scholarships_awards', 'community_college', 'major_minor',
    #    'other_universities', 'top_100_universities', 'top_10_universities',
    #    'associate_education_level', 'bachelor_education_level',
    #    'doctor_education_level', 'master_education_level', 'courses',
    #    'languages', 'hobbies', 'technical_skills', 'company_consumer',
    #    'company_energychem', 'company_fin', 'company_foodbev',
    #    'company_health', 'company_industrial', 'company_other',
    #    'company_services', 'company_tech', 'HonorSociety', 'LatinHonors',
    #    'ScholarshipsAward', 'CommCollege', 'OtherUni', 'Top100Uni', 'Top10Uni',
    #    'Associates', 'Bachelors', 'Masters', 'Doctors', 'FoodBev', 'Consumer',
    #    'EnergyChem', 'Fin', 'HealthMed', 'Industrial', 'Tech', 'Services',
    #    'OtherCo', 'CFA', 'CQE', 'HACCP', 'CPA', 'ACCA', 'CMA', 'CIMA', 'CFCS',
    #    'CPIM', 'PHR', 'PMP', 'FEexam', 'LeanSixSigmaBlackBelt', 'CHCM',
    #    'OSHACertified', 'CSP', 'ASP', 'CertifiedSanitarian', 'CBA', 'CCT',
    #    'CFSQA', 'CMQOE', 'CMBB', 'CQA', 'CQIA', 'CQI', 'CQPA', 'CQT', 'CSSBB',
    #    'CSSGB', 'CSSYB', 'CSQE', 'CSQP', 'Accounting', 'AdultDevelopment',
    #    'Aeronautics', 'AfricanStudies', 'AgriculturalBusiness', 'Agriculture',
    #    'AirwayScience', 'AmericanStudies', 'Analytics', 'AnimalScience',
    #    'Archaeology', 'Architecture', 'ArtStudies', 'ArtificialIntelligence',
    #    'AsianStudies', 'AutomotiveEngineering', 'Aviation', 'BakeryScience',
    #    'Biology', 'BiomedicalEngineering', 'Biophysics', 'BioresourceScience',
    #    'BusinessAdministration', 'BusinessAnalytics', 'ChemicalEngineering',
    #    'ChildCare', 'CivilEngineering', 'CommunicationTechnology',
    #    'Communication', 'ComputerEngineering', 'ComputerNetworks&Systems',
    #    'ComputerScience', 'ConstructionEngineering', 'Counseling',
    #    'CriminalScience', 'DataManagementTechnology', 'DataScience',
    #    'DentalHygiene', 'DieselTechnology', 'DraftingTechnology',
    #    'EarlyChildhoodEducation', 'Ecology', 'Economics',
    #    'EducationAdministration', 'ElectricalEngineering',
    #    'EnglishAsASecondLanguage', 'EnglishEducation', 'EnvironmentalDesign',
    #    'EnvironmentalEngineering', 'EnvironmentalStudies', 'EquestrianStudies',
    #    'EthnicStudies', 'EuropeanStudies', 'FacilitiesAdministration',
    #    'Fashion', 'Film', 'Finance', 'FitnessManagement', 'Food&Nutrition',
    #    'Forensics', 'Forestry', 'FuneralServices', 'GenderStudies', 'Genetics',
    #    'Genomics', 'Geography', 'Geology', 'GlobalHealth', 'GraphicDesign',
    #    'HealthServicesAdministration', 'HispanicStudies', 'HistoryofArt',
    #    'History', 'HomeEconomics', 'Horticulture', 'Hospitality',
    #    'HumanDevelopment', 'HumanResources', 'HumanRightsStudies',
    #    'Immunology', 'IndustrialProductionTechnologies', 'Informatics',
    #    'InformationScience', 'InformationTechnology', 'InteriorDesign',
    #    'InternationalBusiness', 'InternationalRelations', 'Journalism',
    #    'Kinesiology', 'LaborRelations', 'LandscapeArchitecture',
    #    'LawEnforcement', 'LawPolitics&Society', 'Law', 'LegalAssistant',
    #    'LiberalStudies', 'LiteraryStudies', 'Literature', 'Logistics',
    #    'Management', 'MarineScience', 'Marketing', 'MaterialManagement',
    #    'MaterialsScience', 'MathematicsEducation', 'Mathematics',
    #    'MechanicalEngineering', 'MedicalAssistant', 'MedicalBilling',
    #    'MedicalImagingScience', 'MedicalTechnology', 'MentalHealthServices',
    #    'Merchandising', 'Meteorology', 'MiddleEasternStudies',
    #    'MiddleSchoolEducation', 'MilitaryTechnologies', 'Music',
    #    'Neuroscience', 'Nursing', 'Oceanography', 'OfficeManagement',
    #    'OrganizationalBehavior', 'OsteopathicMedicine', 'PackagingEngineering',
    #    'Parks&Recreation', 'Pathology', 'Pharmacy', 'Philosophy',
    #    'PhysicalScience', 'Physics', 'PoliticalScience', 'Psychology',
    #    'PublicHealth', 'PublicPolicy', 'PublicRelations',
    #    'PurchasingManagement', 'QuantumComputing', 'QuantumMechanics',
    #    'Radiology', 'RehabilitationTherapy', 'ResourceManagement',
    #    'RetailManagement', 'RiskManagement', 'SeniorHighEducation',
    #    'SocialScienceEducation', 'SocialScience', 'SocialWelfare', 'Sociology',
    #    'SpeechStudies', 'SportManagement', 'SubstanceAbuseCounseling',
    #    'Surveying', 'SystemsEngineering', 'TheaterStudies', 'Theology',
    #    'Tourism', 'UrbanPlanning', 'UrbanStudies', 'VeterinarianAssisting',
    #    'VeterinaryScience', 'VisualCommunication', 'Viticulture', 'WarStudies',
    #    'WebDesign', 'WebDevelopment', 'WeldingTechnology',
    #    'WildlifeManagement', 'WorldLanguages', 'Zoology', 'Advertising',
    #    'AgriculturalEducation', 'Anatomy', 'ArtEducation', 'Biochemistry',
    #    'Bioengineering', 'Bioethics', 'BiologicalEngineering', 'Biotechnology',
    #    'Business', 'BusinessEducation', 'CeramicEngineering', 'Chemistry',
    #    'Chiropractic', 'ClinicalTrialsManagement', 'CognitiveScience',
    #    'CommunityAdvocacy', 'CommunityOrganization', 'ConflictResolution',
    #    'ContractsManagement', 'CyberSecurity', 'Dentistry', 'Design',
    #    'EducationalLeadership', 'EducationalTechnology', 'Endocrinology',
    #    'EngineeringScience', 'Entrepreneurship', 'Epidemiology',
    #    'Floriculture', 'FoodServicesManagement', 'Genealogy', 'GlobalStudies',
    #    'HealthEducation', 'HistoricPreservation', 'HomelandSecurity',
    #    'HumanServices', 'IndustrialDesign', 'IndustrialEngineering',
    #    'IndustrialManagement', 'JuridicalScience', 'LibraryScience',
    #    'ManagerialEconomics', 'Medicine', 'MetallurgicalEngineering',
    #    'MineralEngineering', 'MinorityStudies', 'MolecularEngineering',
    #    'NuclearEngineering', 'OccupationalSafety', 'Optometry', 'Paleontology',
    #    'PetroleumEngineering', 'Pharmacogenomics', 'Phlebotomy',
    #    'PhysicalEducation', 'PhysicianAssistant', 'PostsecondaryEducation',
    #    'ProjectManagement', 'PublicAdministration', 'QualityControl',
    #    'RealEstate', 'RespiratoryTherapy', 'SafetyTechnologies',
    #    'ScienceEducation', 'SecretarialStudies', 'SoftwareEngineering',
    #    'SoundEngineering', 'SpecialEducation', 'Statistics',
    #    'TechnologyEducation', 'Thermodynamics', 'Toxicology', 'Transportation',
    #    'UrbanTeacherEducation', 'Akkadian', 'Anatolian', 'Arabic', 'Aramaic',
    #    'Armenian', 'Bengali', 'Basque', 'Bosnian', 'Cantonese', 'Catalan',
    #    'Chinese', 'Croatian', 'Czech', 'Danish', 'Dutch', 'Egyptian',
    #    'English', 'French', 'Geez', 'German', 'Greek', 'Hebrew', 'Hindi',
    #    'Hokkien', 'Italian', 'Japanese', 'Kazakh', 'Korean', 'Latin', 'Malay',
    #    'Mandarin', 'Marathi', 'Norwegian', 'Pali', 'Persian', 'Polish',
    #    'Portuguese', 'Punjabi', 'Russian', 'Sanskrit', 'Serbian',
    #    'SignLanguange', 'Slavic', 'Spanish', 'Sumerian', 'Swahili', 'Tamil',
    #    'Telugu', 'Tibetan', 'Turkish', 'Welsh', 'Urdu', 'Uzbek', 'Yiddish',
    #    'APstyle', 'Abletonlive', 'Abaqus', 'AcrobatDistiller',
    #    'AdobeFireworks', 'AdobeFlash', 'Airflow', 'Algorithms',
    #    'AppDevelopment', 'Ariba', 'AutoCAD', 'Axioma', 'Azure', 'Azkaban',
    #    'Adobecreativesuite', 'Agilemethodology', 'AmazonWebServices',
    #    'ApplicationProgrammingInterface', 'Bigdata', 'BirdEye', 'Bloomberg',
    #    'C#', 'C[+][+]', 'Ceridian', 'CREST', 'CRD', 'CSS', 'Clustering',
    #    'CognosAnalytics', 'ContentManagementSystems', 'DBMS', 'DHTML',
    #    'Databasedesign', 'Datapipeline', 'Dataarchitecture', 'Dataengineering',
    #    'Datamining', 'Decisiontrees', 'Deltek', 'Docker', 'Dreamweaver',
    #    'Datavisualization', 'Elixir', 'Excel', 'EPARegulation', 'FactSet',
    #    'Facebookforbusiness', 'FinalCut', 'FSSC22000', 'Googleadwords',
    #    'Googleanalytics', 'GraphicUserInterface', 'H20', 'HRIS',
    #    'Hubspotsales', 'Hyperion', 'Hadoop', 'Hbase', 'IBMCloud',
    #    'IECstandards', 'ISOstandards', 'JQuery', 'JSON', 'Jenkins', 'Jsp',
    #    'Jspf', 'JavaScript', 'Java', 'Kanban', 'Keras', 'Keywordoptimization',
    #    'Kronos', 'LaTeX', 'LexisNexis', 'Linux', 'Luigi', 'MATLAB', 'MYOB',
    #    'Matlab', 'Machinelearning', 'Mathematica', 'Maven', 'Mendeley',
    #    'MongoDB', 'MicrosoftAccess', 'MicrosoftPowerPoint',
    #    'MicrosoftPublisher', 'MicrosoftWord', 'Microsoftoffice', 'NVivo',
    #    'Neo4j', 'Neuralnetworks', 'NodeXL', 'ObjectiveC', 'Oracle',
    #    'ORMregulation', 'PCF', 'PHP', 'Peoplesoft', 'PowerBI', 'Precima',
    #    'PRINCE2', 'Python', 'Pylearn2', 'PostgreSQL', 'Qualtrics',
    #    'Quickbooks', 'QuarkXpress', 'RAPTOR', 'RoboticProcessAutomation',
    #    'Rust', 'Ruby', 'sarbanesoxley', 'SDLC', 'ServiceNow', 'SPSS', 'SQL',
    #    'SQLite', 'STATA', 'Scala', 'Sendgrid', 'SmartCo', 'Smalltalk',
    #    'SolidWorks', 'Softwaredevelopment', 'SoftwareDevelopment',
    #    'SparkStreaming', 'Storm', 'SupplyTrack', 'Surveyxact', 'Swift',
    #    'Syteline', 'SQLServer', 'Searchengineoptimization', 'Sklearn', 'Spark',
    #    'Tableau', 'Typeform', 'Tensorflow', 'UIdesign', 'UWP', 'Ultipro',
    #    'Unittesting', 'Unix', 'VHDL', 'VBA', 'Visio', 'WML', 'WPF',
    #    'Waterfall', 'Weblogic', 'WestLaw', 'Webscraping', 'WindowsServer',
    #    'Windows', 'Wordpress', 'XAML', 'XHTML', 'XML', 'AWK', 'SEB', 'Sed',
    #    'C', 'SAP', 'R', 'SAS', 'Hive', 'Html', 'SSIS', 'GCP', 'Git', 'Mac',
    #    'operatingsystems', 'StemMajor', 'GPAmax']
    #
    # job_dummies_ideal = job_dummies_ideal[cols]

    job_features_ideal = job_dummies_ideal[job_dummies_ideal.columns[:44]]
    resume_features = resume_dummies[resume_dummies.columns[:47]]
    job_dummies_ideal.drop(job_dummies_ideal.columns[1:44], axis=1, inplace=True)
    resume_dummies.drop(resume_dummies.columns[2:47], axis=1, inplace=True)

    # read raw df text data and vectorize for embeddings
    EMAIL_REGEX = r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,4}"

    def remove_stop_words(x):
        try:
            x = re.sub('(\w)(/)(\w)',r'\1 \3', x)  # for responsibilities/accountabilities
            x = re.sub(EMAIL_REGEX, "", x)  # for email
            x = re.sub('[^a-z- ]', '', x)  # for punctuation
            x = re.sub(' - ', ' ', x)  # for dashes
            x = re.sub('-', ' ', x)  # for dashes
            x = re.sub(' \w ', ' ', x)  # for single letters
            x = re.sub('\s\s+', ' ', x)  # for multiple spaces

            word_tokens = word_tokenize(x)
            filtered_sentence = [w for w in word_tokens if not w in STOP_WORDS]
            return filtered_sentence
        except Exception as e:
            print("error:", e)
            print('string:', x)
            sys.exit(1)

    job_text = job_features_ideal[['ReqID', 'text']]
    job_text["text"].replace(r'[\d]', '', regex=True, inplace=True)
    job_text["text"] = job_text["text"].astype(str).apply(lambda x: x.lower().replace('\r', ' ').replace('\n', ' ').replace('\t', ' '))
    job_text["text"] = job_text["text"].apply(remove_stop_words)

    resume_text = resume_features[['ReqID', 'CanID', 'text']]
    resume_text["text"].replace(r'[\d]', '', regex=True, inplace=True) # remove numbers
    resume_text["text"] = resume_text["text"].astype(str).apply(lambda x: x.lower().replace('\r', ' ').replace('\n', ' ').replace('\t', ' '))
    # removing 1 and 2 letter words
    resume_text["text"] = resume_text["text"].apply(lambda x: re.sub('\s\w{1,2}\s', ' ', x))
    resume_text["text"] = resume_text["text"].apply(remove_stop_words)

    resume_dummies.rename(columns={'CanID': 'ID'}, inplace=True)
    job_dummies_ideal['ID'] = job_dummies_ideal['ReqID']
    job_dummies_ideal = job_dummies_ideal[resume_dummies.columns]
    all_dummies_ideal = pd.concat([resume_dummies, job_dummies_ideal])

    def GenerateCountEmbedding(req_id, job_text_df, resume_text_df, all_resumes):
        pos_jd_text = job_text_df[job_text_df["ReqID"] == req_id]
        if all_resumes:
            pos_resume_text = resume_text_df
            pos_resume_text.drop_duplicates(keep='first', inplace=True, subset='CanID')
        else:
            pos_resume_text = resume_text_df[resume_text_df["ReqID"] == req_id]
        pos_jd_text.rename(columns={'ReqID': 'ID', 'Job Description': 'text'}, inplace=True)
        pos_jd_text.ID = req_id
        pos_jd_text = pos_jd_text[['ID', 'text']]
        pos_resume_text.rename(columns={'CanID': 'ID', 'Resume Text': 'text'}, inplace=True)
        pos_resume_text = pos_resume_text[['ID', 'text']]

        df = pos_jd_text.append(pos_resume_text)
        df.set_index('ID', inplace=True)

        df['text'] = df['text'].apply(lambda x: [''] if x[0] == 'nan' else x)
        df['text'] = df['text'].apply(lambda x: ' '.join(x))
        count = CountVectorizer()
        pos_embedding = count.fit_transform(df['text'])
        pos_embedding = pd.DataFrame(pos_embedding.toarray())
        pos_embedding.insert(loc=0, column="ID", value=df.index)

        return pos_embedding

    def RecommendTop(jobID, full_df):  # Cos Sim and rank candidates
        # returns x recommended resume ID's based on Job Description
        recommended_candidates = []
        full_df.fillna(0, inplace=True)
        indices = pd.Series(full_df["ID"])
        cos_sim = cosine_similarity(full_df.drop("ID", axis=1))  # pairwise similarities for all samples in the df
        try:
            idx = indices[indices == jobID].index[0]
            score_series = pd.Series(cos_sim[idx]).sort_values(ascending=False)
            for i in score_series.index:
                recommended_candidates.append(list(indices)[i])
        except IndexError:
            print(jobID, 'had and index error')
        output = pd.DataFrame({'Candidate ID': recommended_candidates, 'cosine': score_series})
        output = output[output['Candidate ID'] != jobID]
        return output

    count_embeddings = GenerateCountEmbedding(jobID, job_text, resume_text, all_resumes)
    count_embeddings['ReqID'] = np.repeat(jobID, len(count_embeddings))
    all_features = pd.DataFrame(count_embeddings).merge(all_dummies_ideal, how="left", on=["ID", 'ReqID'])
    rankings = RecommendTop(jobID=jobID, full_df=all_features.drop('ReqID', axis=1))

    rankings = rankings[:topX]
    rankings.reset_index(drop=True, inplace=True)

    return rankings, job_features_ideal.text[job_features_ideal.ReqID == jobID].values, all_features


import numpy as np
from . import scene, cuda, ragged
import torch
import logging
from rebar import arrdict, dotdict
import matplotlib.pyplot as plt

log = logging.getLogger(__name__)

AGENT_WIDTH = .15
TEXTURE_RES = .05

# Used for collision radius and near camera plane
AGENT_RADIUS = 1/2**.5*AGENT_WIDTH

def gamma_encode(x):
    """Converts RGB data to `viewable values <https://en.wikipedia.org/wiki/Gamma_correction>`_."""
    return x**(1/2.2)

def gamma_decode(x):
    """Converts RGB data to `interpolatable values <https://en.wikipedia.org/wiki/Gamma_correction>`_."""
    return x**2.2

def _init_agents(n_envs, n_agents, device='cuda'):
    """Creates and returns an Agents datastructure"""
    data = arrdict.arrdict(
            angles=torch.zeros((n_envs, n_agents)),
            positions=torch.zeros((n_envs, n_agents, 2)),
            angvelocity=torch.zeros((n_envs, n_agents)),
            velocity=torch.zeros((n_envs, n_agents, 2)))
    return cuda.Agents(**data.to(device))

class Core:

    def __init__(self, scenery, res=64, fov=130, fps=10):
        """The core rendering and physics interface.

        To create the Core, you pass a :class:`~megastep.cuda.Scenery` that describes the environment. Once created,
        the tensors hanging off of the Core give the state of the world, and that state can be advanced with the
        functions in :mod:`~megastep.cuda`.

        :var agents: A :class:`~megastep.cuda.Agents` object describing the agents.
        :var scenery: A :class:`~megastep.cuda.Scenery` object describing the scene.
        :var progress: A (n_env, n_agent)-tensor giving how far the agent was able to move in the previous timestep as a
            fraction of its intended movement, before running into an obstable. A value less than 1 means the agent collided
            with something. Useful for detecting collisions.
        :var n_envs: Number of environments. Same as the number of geometries passed in.
        :var n_agents: Number of agents.
        :var res: The horizontal resolution of observations.
        :var fov: The field of view in degrees.
        :var agent_radius: The radius of a disc containing the agent, in meters.
        :var fps: The framerate.
        :var random: The seeded :class:`numpy.random.RandomState` used to initialize the environment. By reusing this
            in any extra random decisions made when generating the environments, it can be guaranteed you'll get the same
            environments every time.

        :param scenery: Describes the static parts of the environment.
        :type scenery: :class:`~megastep.cuda.Scenery`.
        :param n_agents: the number of agents to put in each environment. Defaults to 1.
        :type n_agents: int
        :param res: The horizontal resolution of the observations. The resolution must be less than 1024, as
            that's the `maximum number of CUDA threads in a block
            <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities>`_. Defaults to 64
            pixels.
        :type res: int
        :param fov: The field of view in degrees. Must be less than 180° due to how frames are rendered. Defaults to 130°.
        :type fov float:
        :param fps: The simulation frame rate/step rate. Defaults to 10.
        :type fps: int
        """

        self.n_envs = len(scenery.lines.widths)
        self.n_agents = scenery.n_agents
        self.res = res
        self.fov = fov
        # TODO: Make this adjustable
        self.agent_radius = AGENT_RADIUS
        self.fps = fps
        self.random = np.random.RandomState(1)

        # TODO: This needs to be propagated to the C++ side
        self.device = scenery.model.device

        assert fov < 180, 'FOV should be less than 180°'

        cuda.initialize(self.agent_radius, self.res, self.fov, self.fps)
        self.scenery = scenery
        self.agents = _init_agents(self.n_envs, self.n_agents, self.device)
        self.progress = torch.ones((self.n_envs, self.n_agents), device=self.device)

    def state(self, e):
        """Returns a :class:`~rebar.dotdict.dotdict` tree representing the state of environment ``e``.

        A typical state looks like this::

            arrdict:
            n_envs          1
            n_agents        4
            res             512
            fov             60
            agent_radius    0.10606601717798211
            fps             10
            scenery           arrdict:
                            model       Tensor((8, 2, 2), torch.float32)
                            lines       Tensor((307, 2, 2), torch.float32)
                            lights      Tensor((21, 3), torch.float32)
                            textures    <megastepcuda.Ragged2D object at 0x7fba34112eb0>
                            baked       <megastepcuda.Ragged1D object at 0x7fba34112670>
            agents          arrdict:
                            angles       Tensor((4,), torch.float32)
                            positions    Tensor((4, 2), torch.float32)
            progress        Tensor((4,), torch.float32)

        This state tree is usually passed onto a :ref:`plotting` function."""
        options = ('n_envs', 'n_agents', 'res', 'fov', 'agent_radius', 'fps')
        options = {k: getattr(self, k) for k in options}

        return arrdict.clone(dotdict.dotdict(
                **options,
                scenery=self.scenery.state(e),
                agents=self.agents.state(e),
                progress=self.progress[e]))

    @classmethod
    def plot_state(cls, state, ax=None, zoom=False):
        from . import plotting
        ax = ax or plt.axes()
        plotting.plot_lines(ax, state, zoom=zoom)
        plotting.plot_lights(ax,state)
        plotting.adjust_view(ax, state, zoom=zoom)
        plotting.plot_fov(ax, state)
        ax.set_xticks([])
        ax.set_yticks([])
        return ax

    def env_full(self, x):
        """Returns a (n_envs,)-tensor on the environment's device full of `x`.

        This isn't strictly required by the Core, but you find yourself making these vectors so often it's useful sugar.
        """
        dtypes = {bool: torch.bool, int: torch.int32, float: torch.float32}
        return torch.full((self.n_envs,), x, device=self.device, dtype=dtypes[type(x)])

    def agent_full(self, x):
        """Returns a (n_envs, n_agents)-tensor on the environment's device full of `x`.

        This isn't strictly required by the Core, but you find yourself making these vectors so often it's useful sugar.
        """
        dtypes = {bool: torch.bool, int: torch.int32, float: torch.float32}
        return torch.full((self.n_envs, self.n_agents), x, device=self.device, dtype=dtypes[type(x)])


import numpy as np
import cv2 as cv
import sys
import os
from tempfile import TemporaryDirectory

FRAMERATE = 25
PROG_START_THRESHOLD = 180 * FRAMERATE
SHOW_BLOCK_MINIMUM = (3 * 60 + 30) * FRAMERATE
BLANK_NEAR_THRESHOLD = 2 * FRAMERATE
SEG_START_DELAY = int(1.0 * FRAMERATE)
SEG_END_DELAY = int(0.75 * FRAMERATE)

# Set params
in_file = sys.argv[1]
out_files = sys.argv[2:]

# Create a new tempdir
tmpdir = TemporaryDirectory()

# Create an array for the blank frame indexes
blanks = [-BLANK_NEAR_THRESHOLD]

# Collect blank frames
cap = cv.VideoCapture(in_file)
i = 1
while(cap.isOpened()):
    ret, frame = cap.read()
    if frame is None:
        break
    if frame.mean() < 0.5 and (i - blanks[-1]) > BLANK_NEAR_THRESHOLD:
        print('Blank frame at %i' % i)
        blanks.append(i)
    i += 1
cap.release()
cv.destroyAllWindows()

blanks.pop(0)

# Find the program start
program_start = None
for frame in list(blanks):
    if(frame > PROG_START_THRESHOLD):
        break
    program_start = frame
    blanks.pop(0) # Pop the frame to add the last one later

print("Program start: %if" % program_start)

# Strip in-ad blanks
i = 0
in_ad = False
while i < len(blanks):
    frame = blanks[i]
    if in_ad: # In an ad
        #if frame - blanks[i - 1] < AD_BREAK_MINIMUM:
        if i + 1 < len(blanks) and blanks[i + 1] - blanks[i] < SHOW_BLOCK_MINIMUM:
            print("Removing in-ad blank %i" % frame)
            blanks.pop(i)
            # i -= 1 # Wind back the for loop as we just deleted an element
        else:
            in_ad = False
            i += 1
    else:
        in_ad = True
        i += 1

blanks.insert(0, program_start)
print("Blanks: %s" % blanks)

# Split the program segments
print('Splitting program...')
i = 0
blanks = list(zip(blanks[0::2], blanks[1::2]))
for i in range(len(blanks)):
    start_s = (blanks[i][0] + SEG_START_DELAY) / FRAMERATE
    end_s   = (blanks[i][1] + SEG_END_DELAY) / FRAMERATE
    time_s  = end_s - start_s
    print("%i: %i to %i" % (i, blanks[i][0] + SEG_START_DELAY, blanks[i][1] + SEG_END_DELAY))
    print("    %fs to %fs, len %fs" % (start_s, end_s, time_s))
    os.system("ffmpeg -ss %f -i \"%s\" -t %f -c:v copy -c:a copy %s/%i.mkv" % (start_s, in_file, time_s, tmpdir.name, i))

# Compute how many segments belong in each file
len_of_file = int(len(blanks) / len(out_files))

# Compile the output files
for out_file_no in range(len(out_files)):
    out_file = out_files[out_file_no]
    print('Creating file list for %s' % out_file)
    with open(('%s/filelist.txt' % tmpdir.name), 'w') as filelist:
        for i in range(out_file_no * len_of_file, out_file_no * len_of_file + len_of_file):
            filelist.write("file %s.mkv\n" % i)

    print('Compiling segments...')
    os.system("ffmpeg -f concat -i %s/filelist.txt -c:v h264 -crf 18 -c:a aac -strict -2 -async 1 \"%s\"" % (tmpdir.name, out_file))

print('Cleaning up...')
tmpdir.cleanup()


import torch
import cv2
import numpy as np
from skimage import transform

from Retinaface.utils.alignment import get_reference_facial_points, FaceWarpException
from Retinaface.utils.box_utils import decode, decode_landmark, prior_box, nms
from Retinaface.utils.config import cfg_slim, cfg_rfb, cfg_mnet, cfg_re50
from Retinaface.models.retinaface import RetinaFace
from Retinaface.models.rfb import RFB
from Retinaface.models.slim import Slim

class FaceDetector:

    def __init__(self, name, weight_path, device='cpu', confidence_threshold=0.99,
                 top_k=5000, nms_threshold=0.4, keep_top_k=750, face_size=(112, 112)):
        """
        RetinaFace Detector with 5points landmarks
        Args:
            name: name of backbone (resnet, mobilenet, slim, rfb)
            weight_path: path of network weight
            device: running device (cuda, cpu)
            face_size: final face size
            face_padding: padding for bounding boxes
        """

        model, cfg = None, None
        if name == 'mobilenet':
            cfg = cfg_mnet
            model = RetinaFace(cfg=cfg, phase='test')
        elif name == 'resnet':
            cfg = cfg_re50
            model = RetinaFace(cfg=cfg, phase='test')
        elif name == 'slim':
            cfg = cfg_slim
            model = Slim(cfg=cfg_slim, phase='test')
        elif name == 'rfb':
            cfg = cfg_rfb
            model = RFB(cfg=cfg_rfb, phase='test')
        else:
            exit('from FaceDetector Exit: model not support \n just(mobilenet, resnet, slim, rfb)')

        # setting for model
        model.load_state_dict(torch.load(weight_path))
        model.to(device).eval()
        self.model = model
        self.device = device
        self.cfg = cfg
        # setting for face detection
        self.thresh = confidence_threshold
        self.top_k = top_k
        self.nms_thresh = nms_threshold
        self.keep_top_k = keep_top_k
        # setting for face align
        self.trans = transform.SimilarityTransform()
        self.out_size = face_size
        self.ref_pts = get_reference_facial_points(output_size=face_size)
        print('from FaceDetector: weights loaded')
        return

    def preprocessor(self, img_raw):
        img = torch.tensor(img_raw, dtype=torch.float32).to(self.device)
        scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]).to(self.device)
        img -= torch.tensor([104, 117, 123]).to(self.device)
        img = img.permute(2, 0, 1).unsqueeze(0)
        return img, scale

    def detect_faces(self, img_raw):
        """
        get a image from ndarray, detect faces in image
        Args:
            img_raw: original image from cv2(BGR) or PIL(RGB)
        Notes:
            coordinate is corresponding to original image
            and type of return image is corresponding to input(cv2, PIL)
        Returns:
            boxes:
                faces bounding box for each face
            scores:
                percentage of each face
            landmarks:
                faces landmarks for each face
        """

        img, scale = self.preprocessor(img_raw)
        # tic = time.time()
        with torch.no_grad():
            loc, conf, landmarks = self.model(img)  # forward pass
            # print('net forward time: {:.4f}'.format(time.time() - tic))

        priors = prior_box(self.cfg, image_size=img.shape[2:]).to(self.device)
        boxes = decode(loc.data.squeeze(0), priors, self.cfg['variance'])
        boxes = boxes * scale
        scores = conf.squeeze(0)[:, 1]
        landmarks = decode_landmark(landmarks.squeeze(0), priors, self.cfg['variance'])
        scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                               img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                               img.shape[3], img.shape[2]]).to(self.device)
        landmarks = landmarks * scale1

        # ignore low scores
        index = torch.where(scores > self.thresh)[0]
        boxes = boxes[index]
        landmarks = landmarks[index]
        scores = scores[index]

        # keep top-K before NMS
        order = scores.argsort(dim=0, descending=True)[:self.top_k]
        boxes = boxes[order]
        landmarks = landmarks[order]
        scores = scores[order]

        # Do NMS
        keep = nms(boxes, scores, self.nms_thresh)
        boxes = torch.abs(boxes[keep, :])
        scores = scores[:, None][keep, :]
        landmarks = landmarks[keep, :].reshape(-1, 5, 2)

        # # keep top-K faster NMS
        landmarks = landmarks[:self.keep_top_k, :]
        scores = scores[:self.keep_top_k, :]
        boxes = boxes[:self.keep_top_k, :]

        return boxes, scores, landmarks

    def detect_align(self, img):
        """
        get a image from ndarray, detect faces in image,
        cropped face and align face
        Args:
            img: original image from cv2(BGR) or PIL(RGB)
        Notes:
            coordinate is corresponding to original image
            and type of return image is corresponding to input(cv2, PIL)

        Returns:
            faces:
                a tensor(n, 112, 112, 3) of faces that aligned
            boxes:
                face bounding box for each face
            landmarks:
                face landmarks for each face
        """
        boxes, scores, landmarks = self.detect_faces(img)

        warped = []
        for src_pts in landmarks:
            if max(src_pts.shape) < 3 or min(src_pts.shape) != 2:
                raise FaceWarpException('facial_pts.shape must be (K,2) or (2,K) and K>2')

            if src_pts.shape[0] == 2:
                src_pts = src_pts.T

            if src_pts.shape != self.ref_pts.shape:
                raise FaceWarpException('facial_pts and reference_pts must have the same shape')

            self.trans.estimate(src_pts.cpu().numpy(), self.ref_pts)
            face_img = cv2.warpAffine(img, self.trans.params[0:2, :], self.out_size)
            warped.append(face_img)

        faces = torch.tensor(warped).to(self.device)
        return faces, boxes, scores, landmarks


import math
import scipy
import scipy.misc
import torch
import numpy as np
from PIL import Image
from lib.models import get_eye_alignment_net, get_cls_net
from lib.config import config, config_imagenet, merge_configs

def get_model_by_name(checkpoint_path, config_path, model_type='landmarks', device='cuda'):

    merge_configs(config, config_path)

    config.defrost()
    config.MODEL.INIT_WEIGHTS = False
    config.freeze()

    if model_type == 'landmarks':
        model = get_eye_alignment_net(config)
    else:
        model = get_cls_net(config_imagenet)

    state_dict = torch.load(checkpoint_path)
    if 'state_dict' in state_dict.keys():
        state_dict = state_dict['state_dict']
        model.load_state_dict(state_dict)
    else:
        model.module.load_state_dict(state_dict)
    model.eval()
    model.to(device)

    return model

def get_lmks_by_img(model, img, output_size=(160, 96), rot=0, conf_score=False):
#   img = np.array(Image.open(image_path).convert('RGB'), dtype=np.float32)

    face_center = torch.Tensor([img.shape[1]//2, img.shape[0]/2])      # (96,96)
    crop_scale = max((img.shape[1]) / output_size[0], (img.shape[0]) / output_size[1])  # 아웃풋과 이미지와의 스케일 비율
    # hegith 와 width 중 더 큰 스케일 차이  --> 192/192 = 4

    img_crop = crop(img, face_center, crop_scale, output_size=output_size, rot=rot)   # face_center : (96,96)
    # 원래 이미지 형태는 행, 열 , 채널  --> 세로 , 가로 , 채널 (0 ,1 ,2)
    img_crop = (img_crop/255.0 - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225]) # normalize ?
    img_crop = img_crop.transpose([2, 0, 1]) # Tensor의 형태로 맞춰 주기 위함  (세로 가로 채널 ) --> (채널 , 세로 , 가로)
    img_crop = torch.tensor(img_crop, dtype=torch.float32).unsqueeze(0).cuda() # batch size의 크기를 넣어줌
    with torch.no_grad():
        pred = model(img_crop)
    return decode_preds(pred, [face_center], [crop_scale], [output_size[0]/4,output_size[1]/4], conf_score)
    #  decode_preds( 히트맵 , [[96,96]], [1], res=[48,48] , conf_score = Flase ) 히트맵의 위치 정보만 필요로 함.

def get_lmks_by_img_v2(model, img, output_size=(96, 160), rot=0, conf_score=False):
#   img = np.array(Image.open(image_path).convert('RGB'), dtype=np.float32)

    face_center = torch.Tensor([img.shape[1]//2, img.shape[0]/2])      # (96,96)
    crop_scale = max((img.shape[1]) / output_size[0], (img.shape[0]) / output_size[1])  # 아웃풋과 이미지와의 스케일 비율
    # hegith 와 width 중 더 큰 스케일 차이  --> 192/192 = 4

    img_crop = crop(img, face_center, crop_scale, output_size=output_size, rot=rot)   # face_center : (96,96)
    # 원래 이미지 형태는 행, 열 , 채널  --> 세로 , 가로 , 채널 (0 ,1 ,2)
    # img_crop = (img_crop/255.0 - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225]) # normalize ?
    img_crop = (img_crop/255.0 - np.array([0.3992,0.3992,0.3992])) / np.array([0.1981,0.1981,0.1981])
    img_crop = img_crop.transpose([2, 0, 1]) # Tensor의 형태로 맞춰 주기 위함  (세로 가로 채널 ) --> (채널 , 세로 , 가로)
    img_crop = torch.tensor(img_crop, dtype=torch.float32).unsqueeze(0).cuda() # batch size의 크기를 넣어줌
    with torch.no_grad():
        pred , pred_gaze = model(img_crop)
    return decode_preds(pred, [face_center], [crop_scale], [output_size[0]/4,output_size[1]/4], conf_score) , pred_gaze
    #  decode_preds( 히트맵 , [[96,96]], [1], res=[48,48] , conf_score = Flase ) 히트맵의 위치 정보만 필요로 함.

def get_preds(scores, conf_score=False):  # 결과 값이 heatmap에서 각 키포인트를 의미하는 각 채널에서의 가장 큰 값과 그 위치를 알아 보는 함수
    """
    get predictions from score maps in torch Tensor
    return type: torch.LongTensor
    """
    assert scores.dim() == 4, 'Score maps should be 4-dim'   #  output의 차원이 4차원이여야함 ( N x C x H x W ) --> (배치사이즈 , 채널, 행수(세로길이) , 열수(가로길이))
    maxval, idx = torch.max(scores.view(scores.size(0), scores.size(1), -1), 2)  # 이미지를 flatten 하게 피고 나서, 각 채널에서 가장 큰 값을 뽑아내는 듯 --> 즉 키포인트의 위치..
    # maxval shape --> [1,50] idx shape --> [1,50]
    maxval = maxval.view(scores.size(0), scores.size(1), 1) # maxval shape --> [1,50,1]
    idx = idx.view(scores.size(0), scores.size(1), 1) + 1 # idx shape --> [1,50,1] , 그리고 각 값에 +1 해줌

    preds = idx.repeat(1, 1, 2).float() # preds shape --> [1,50,2]

    preds[:, :, 0] = (preds[:, :, 0] - 1) % scores.size(3) + 1 # 행에서 몇칸 가야 하는지
    preds[:, :, 1] = torch.floor((preds[:, :, 1] - 1) / scores.size(3)) + 1 # 몇번째 행인지 확인

    # preds는 output에서 각 채널에 대해 가장 큰 값을 가지는 인덱스를 (열 , 행) 으로 가지고 있음

    pred_mask = maxval.gt(0).repeat(1, 1, 2).float() # pred_mask는 maxval이 0보다 작지 않다면 다 1로 변환해서 (1,50,2) 의 형태로 가지고 있음
    preds *= pred_mask # 각 maxval를 가지는 위치에 마스크 값을 곱해줌 (여기서 0보다 작은 값들을 골라 내주는 거 같음)

    if conf_score == True:
        return preds, maxval  # 결과적으로 max 값을 가지는 위치값과 , max값을 결과값을 줌  preds의 shape는 [1,50,2] --> 히트맵에서 어떤 위치에 같을 가지고 있는지
    else:                                                                             #   max의 shape는   [1,50,1] --> 각 히트맵에서의 최대 값 --> 키 포인트 값을 의미
        return preds

def decode_preds(output, center, scale, res, conf_score=False): # output size에 맞게 결과 값을 얻음
    if conf_score == True:
        coords, scores = get_preds(output, conf_score)  # float type  [1,50,48,48] heat map에서 각 채널 당 가장 큰 값과 그 위치(열 행)를 반 환
    else:
        coords = get_preds(output)  # float type
    coords = coords.cpu()
    # pose-processing
    for n in range(coords.size(0)): # 배치 사이즈 만큼
        for p in range(coords.size(1)): # 각 채널 당
            hm = output[n][p]        # htmap의 모든 정보임
            px = int(math.floor(coords[n][p][0])) # 소수점 버려 버림 (열)
            py = int(math.floor(coords[n][p][1])) # 소수점 버려 버림 (행)
            if (px > 1) and (px < res[0]) and (py > 1) and (py < res[1]): # --> 이 부분이 이해가 안감
                diff = torch.Tensor([hm[py - 1][px] - hm[py - 1][px - 2], hm[py][px - 1]-hm[py - 2][px - 1]])
                coords[n][p] += diff.sign() * .25
    coords += 0.5 # 왜 더하지????
    preds = coords.clone()
    # Transform back
    for i in range(coords.size(0)): # 열 행 순
        preds[i] = transform_preds(coords[i], center[i], scale[i], res)

    if preds.dim() < 3:
        preds = preds.view(1, preds.size())

    if conf_score == True:
        return preds.cpu().numpy().squeeze(0), scores.cpu().numpy().squeeze(0)
    else:
        return preds.cpu().numpy().squeeze(0)

def crop(img, center, scale, output_size=(96,160), rot=0):
    # center : [center_w, center_h]
    center_new = center.clone()

    # Preprocessing for efficient cropping
    ht, wd = img.shape[0], img.shape[1] # height , widht
    sf = scale * 200.0 / output_size[0] # sf는 scale = 4 일 것이다. 200 / 192 --> 약 1
    if sf < 2:
        sf = 1
    else:
        new_size = int(np.math.floor(max(ht, wd) / sf))  # 46?
        new_ht = int(np.math.floor(ht / sf))
        new_wd = int(np.math.floor(wd / sf))
        if new_size < 2:
            return torch.zeros(output_size[0], output_size[1], img.shape[2]) \
                        if len(img.shape) > 2 else torch.zeros(output_size[0], output_size[1])
        else:
            img = np.array(Image.fromarray(img.astype(np.uint8)).resize((new_wd, new_ht)))
#             img = scipy.misc.imresize(img, [new_ht, new_wd])  # (0-1)-->(0-255)
            center_new[0] = center_new[0] * 1.0 / sf
            center_new[1] = center_new[1] * 1.0 / sf
            scale = scale / sf

    # Upper left point
    ul = np.array(transform_pixel([0, 0], center_new, scale, output_size, invert=1))
    # Bottom right point
    br = np.array(transform_pixel(output_size, center_new, scale, output_size, invert=1))

    # Padding so that when rotated proper amount of context is included
    pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2)
    if not rot == 0:
        ul -= pad
        br += pad

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]

    new_img = np.zeros(new_shape, dtype=np.float32)

    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1]
    # Range to sample from original image
    old_x = max(0, ul[0]), min(len(img[0]), br[0])
    old_y = max(0, ul[1]), min(len(img), br[1])
    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1], old_x[0]:old_x[1]]

    if not rot == 0:
        # Remove padding
        new_img = scipy.misc.imrotate(new_img, rot)
        new_img = new_img[pad:-pad, pad:-pad]
    new_img = np.array(Image.fromarray(new_img.astype(np.uint8)).resize(output_size[::-1]))
#     new_img = scipy.misc.imresize(new_img, output_size)
    return new_img

def transform_pixel(pt, center, scale, output_size, invert=0, rot=0):
    # Transform pixel location to different reference
    t = get_transform(center, scale, output_size, rot=rot) # 픽셀의 위취를 다른 reference로 변환
    if invert:
        t = np.linalg.inv(t)
    new_pt = np.array([pt[0] - 1, pt[1] - 1, 1.]).T
    new_pt = np.dot(t, new_pt)
    return new_pt[:2].astype(int) + 1

def get_transform(center, scale, output_size, rot=0):
    """
    General image processing functions
    """
    # Generate transformation matrix
    h = 200 * scale # 스케일을 왜 2배 늘리는거?  h = 800 ( scale = 4)
    t = np.zeros((3, 3))  # 3x3 행렬 만듬
    t[0, 0] = float(output_size[1]) / h                              #
    t[1, 1] = float(output_size[0]) / h                              #  | 192/200    0      192 * (96/200 + 0.5) |
    t[0, 2] = output_size[1] * (-float(center[0]) / h + .5)          #  |   0     192/200   192 * (96/200 + 0.5) |
    t[1, 2] = output_size[0] * (-float(center[1]) / h + .5)          #  |   0        0              1            |
    t[2, 2] = 1
    if not rot == 0: # rotation 진행하는 거 같음
        rot = -rot  # To match direction of rotation from cropping
        rot_mat = np.zeros((3, 3))
        rot_rad = rot * np.pi / 180
        sn, cs = np.sin(rot_rad), np.cos(rot_rad)
        rot_mat[0, :2] = [cs, -sn]
        rot_mat[1, :2] = [sn, cs]
        rot_mat[2, 2] = 1
        # Need to rotate around center
        t_mat = np.eye(3)
        t_mat[0, 2] = -output_size[1]/2
        t_mat[1, 2] = -output_size[0]/2
        t_inv = t_mat.copy()
        t_inv[:2, 2] *= -1
        t = np.dot(t_inv, np.dot(rot_mat, np.dot(t_mat, t)))
    return t

def transform_preds(coords, center, scale, output_size):

    for p in range(coords.size(0)): #
        coords[p, 0:2] = torch.tensor(transform_pixel(coords[p, 0:2], center, scale, output_size, 1, 0))
    return coords


import numpy as np
from PyQt5.QtCore import Qt,QByteArray, QIODevice, QThread, QObject, pyqtSignal,pyqtSlot,QIODevice
from PyQt5.QtMultimedia import QAudioFormat, QAudioOutput
from PyQt5.QtWidgets import QApplication, QWidget,QSlider,QPushButton, QMessageBox, QVBoxLayout,QHBoxLayout,QLabel
from PyQt5.QtGui import QPixmap
import sys
import mido
import random
import matplotlib.pyplot as pl
import Resonator as ResC

SAMPLE_RATE=32000
AUDIO_CHANS=1
SAMPLE_SIZE=16
CTRL_INTERVAL=100 #milliseconds of audio

class MidiPortReader(QObject):
    noteOn=pyqtSignal(int)
    noteVelocity=pyqtSignal(int)
    noteOff=pyqtSignal(int)
    def __init__(self):
        QObject.__init__(self)
        mido.set_backend(
            'mido.backends.rtmidi/LINUX_ALSA'
        )
        print(
            "Using MIDI APIs: {}".format(
                mido.backend.module.get_api_names()
            )
        )
        self.velocity=0
        self.state=0
    def listener(self):

        with mido.open_input(
            'pipes',
            virtual=True
        ) as mip:
            for mmsg in mip:
                print(mmsg.velocity)
                print(mmsg.type)
                if mmsg.type=='note_on':
                    self.noteOn.emit(mmsg.note)
                elif mmsg.type=='note_off':
                    self.noteOff.emit(mmsg.note)
                self.noteVelocity.emit(mmsg.velocity)

class Generator(QIODevice):

    SAMPLES_PER_READ=1024

    def __init__(self,format,parent=None):
        self.format=format
        QIODevice.__init__(self,parent)
        self.data=QByteArray()
        self.state=0
        self.vel=0
        self.vol=0
        self.q=0
        self.d=0
    def start(self):
        self.open(QIODevice.ReadOnly)

    @pyqtSlot(int)
    def noteOn(self,d):
        self.state=1
        f=(2**((d-69)/12))*440
        print(f)
        self.resonator=ResC.Resonator(f,self.q,SAMPLE_RATE)
        print(self.voices)

    @pyqtSlot(int)
    def noteVelocity(self,v):
        self.vel=v/100

    @pyqtSlot(int)
    def noteOff(self,s):
        print('note if off')
        self.state=0

    @pyqtSlot(int)
    def volSlide(self,volume):
        self.vol=volume/100

    @pyqtSlot(int)
    def qFactor(self,q):
        self.q=q
        print(q)

    def generate(self,NO_OF_SAMPLES):
        tone=np.zeros(NO_OF_SAMPLES)
        for i in range(len(tone)):
            tone[i]=self.resonator.filter()
        tone=(self.vel*50*tone*self.state*self.vol).astype(np.int16)
        return tone.tostring()

    def readData(self,bytes):
        if bytes>2*Generator.SAMPLES_PER_READ:
            bytes=2*Generator.SAMPLES_PER_READ
        return self.generate(bytes//2)

class Window(QWidget):
    volSlide=pyqtSignal(int)
    def __init__(self,parent=None):
        #UI
        QWidget.__init__(self,parent)
        self.create_UI(parent)
        #audio formatting
        format=QAudioFormat()
        format.setChannelCount(AUDIO_CHANS)
        format.setSampleRate(SAMPLE_RATE)
        format.setSampleSize(SAMPLE_SIZE)
        format.setCodec("audio/pcm")
        format.setByteOrder(
            QAudioFormat.LittleEndian
        )
        format.setSampleType(
            QAudioFormat.SignedInt
        )
        self.output=QAudioOutput(format,self)
        output_buffer_size=\
            int(2*SAMPLE_RATE \
                *CTRL_INTERVAL/1000)
        self.output.setBufferSize(
            output_buffer_size
        )
        self.generator=Generator(format,self)
        #THREADS
        self.midiListener=MidiPortReader()
        self.listenerThread=QThread()
        self.midiListener.moveToThread(
                self.listenerThread
        )
        self.listenerThread.started.connect(
                self.midiListener.listener
        )
        self.listenerThread.start()
        self.midiListener.noteOff.connect(self.generator.noteOff)
        self.midiListener.noteVelocity.connect(self.generator.noteVelocity)
        self.midiListener.noteOn.connect(self.generator.noteOn)
        self.volumeSlider.valueChanged.connect(self.generator.volSlide)
        self.qfacSlider.valueChanged.connect(self.generator.qFactor)
        self.generator.start()
        self.output.start(self.generator)

    def create_UI(self,parent):
        rockLabel=QLabel()
        rockLabel.setText("Let's ROCK!!!")
        volLabel=QLabel()
        volLabel.setText("Volume")
        self.volumeSlider=QSlider(Qt.Horizontal)
        self.volumeSlider.setMinimum(0)
        self.volumeSlider.setMaximum(100)
        qfacLabel=QLabel()
        qfacLabel.setText("Q-Factor")
        self.qfacSlider=QSlider(Qt.Horizontal)
        self.qfacSlider.setMinimum(0)
        self.qfacSlider.setMaximum(1000)
        self.qfacSlider.setValue(5000)
        self.quitButton=\
                QPushButton(self.tr('&Quit'))
        self.volSlide.emit(7)
        self.lynch=QLabel()
        self.halfDown=QPushButton("+1")
        self.halfUp=QPushButton("-1")
        pixmap=QPixmap()
        pixmap.load('lynch2.jpeg')
        pixmap=pixmap.scaledToWidth(200)
        self.lynch.setPixmap(pixmap)
        vLayout=QVBoxLayout(self)
        h0Layout=QHBoxLayout()
        h0Layout.addWidget(rockLabel)
        h0Layout.addStretch(1)
        h0Layout.addWidget(self.lynch)
        vLayout.addLayout(h0Layout)
        hLayout=QHBoxLayout()
        hLayout.addWidget(volLabel)
        hLayout.addStretch(1)
        hLayout.addWidget(self.volumeSlider)
        vLayout.addLayout(hLayout)
        h2Layout=QHBoxLayout()
        h2Layout.addWidget(qfacLabel)
        h2Layout.addStretch(1)
        h2Layout.addWidget(self.qfacSlider)
        vLayout.addLayout(h2Layout)
        self.quitButton.clicked.connect(
                self.quitClicked
        )

    @pyqtSlot()
    def quitClicked(self):
        self.close()

if __name__== "__main__":
    app=QApplication(sys.argv)
    window=Window()
    window.show()
    sys.exit(app.exec_())