from __future__ import division

import numpy as np
import torch
import torch.nn as nn
from mmcv.cnn import normal_init

from mmdet.core import (AnchorGenerator3D, anchor_target, delta2bbox,
                        multi_apply, weighted_cross_entropy, weighted_smoothl1,
                        weighted_binary_cross_entropy,
                        weighted_sigmoid_focal_loss, multiclass_nms)
from mmdet.core.anchor.anchor_target import get_anchor_inside_flags
from ..registry import HEADS
from mmdet.core import tensor2img3D
import mmcv
import random
import cv2
import torch.nn.functional as F

@HEADS.register_module
class AnchorHead3D(nn.Module):
    """Anchor-based head (RPN, RetinaNet, SSD, etc.).

    Args:
        in_channels (int): Number of channels in the input feature map.
        feat_channels (int): Number of channels of the feature map.
        anchor_scales (Iterable): Anchor scales.
        anchor_ratios (Iterable): Anchor aspect ratios.
        anchor_strides (Iterable): Anchor strides.
        anchor_base_sizes (Iterable): Anchor base sizes.
        target_means (Iterable): Mean values of regression targets.
        target_stds (Iterable): Std values of regression targets.
        use_sigmoid_cls (bool): Whether to use sigmoid loss for classification.
            (softmax by default)
        use_focal_loss (bool): Whether to use focal loss for classification.
    """  # noqa: W605

    def __init__(self,
                 num_classes,
                 in_channels,
                 feat_channels=256,
                 anchor_scales=[8, 16, 32],
                 anchor_depth_scales=[3],
                 anchor_ratios=[0.5, 1.0, 2.0],
                 anchor_strides=[4, 8, 16, 32, 64],
                 anchor_strides_depth=[2, 4, 8, 16, 32],
                 anchor_base_sizes=None,
                 target_means=(.0, .0, .0, .0),
                 target_stds=(1.0, 1.0, 1.0, 1.0),
                 use_sigmoid_cls=False,
                 use_focal_loss=False):
        super(AnchorHead3D, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.feat_channels = feat_channels
        self.anchor_scales = anchor_scales
        self.anchor_depth_scales = anchor_depth_scales
        self.anchor_ratios = anchor_ratios
        self.anchor_strides = anchor_strides
        self.anchor_base_sizes = list(
            anchor_strides) if anchor_base_sizes is None else anchor_base_sizes
        self.anchor_strides_depth = anchor_strides_depth
        self.target_means = target_means
        self.target_stds = target_stds
        self.use_sigmoid_cls = use_sigmoid_cls
        self.use_focal_loss = use_focal_loss
        self.pos_indices = None
        self.pos_indices_test = None

        self.anchor_generators = []
        for anchor_base, anchor_depth_base in zip(self.anchor_base_sizes, self.anchor_strides_depth):
            self.anchor_generators.append(
                AnchorGenerator3D(anchor_base, anchor_scales, anchor_depth_scales, anchor_ratios, anchor_depth_base))

        self.num_anchors = len(self.anchor_ratios) * len(self.anchor_scales)
        if self.use_sigmoid_cls:
            self.cls_out_channels = self.num_classes - 1
        else:
            self.cls_out_channels = self.num_classes

        self._init_layers()

    def forward(self, feats):
        return multi_apply(self.forward_single, feats)

    def get_anchors(self, featmap_sizes, img_metas):
        """Get anchors according to feature map sizes.

        Args:
            featmap_sizes (list[tuple]): Multi-level feature map sizes.
            img_metas (list[dict]): Image meta info.

        Returns:
            tuple: anchors of each image, valid flags of each image
        """
        num_imgs = len(img_metas)
        num_levels = len(featmap_sizes)

        # since feature map sizes of all images are the same, we only compute
        # anchors for one time
        multi_level_anchors = []
        for i in range(num_levels):
            anchors = self.anchor_generators[i].grid_anchors(
                featmap_sizes[i], self.anchor_strides[i], self.anchor_strides_depth[i])
            multi_level_anchors.append(anchors)
        anchor_list = [multi_level_anchors for _ in range(num_imgs)]

        # for each image, we compute valid flags of multi level anchors
        valid_flag_list = []
        for img_id, img_meta in enumerate(img_metas):
            multi_level_flags = []
            for i in range(num_levels):
                anchor_stride = self.anchor_strides[i]
                anchor_depth_stride = self.anchor_strides_depth[i]
                feat_z, feat_h, feat_w = featmap_sizes[i]
                h, w, _, z = img_meta['pad_shape']
                valid_feat_h = min(int(np.ceil(h / anchor_stride)), feat_h)
                valid_feat_w = min(int(np.ceil(w / anchor_stride)), feat_w)
                valid_feat_z = min(int(np.ceil(z / anchor_depth_stride)), feat_z)
                flags = self.anchor_generators[i].valid_flags(
                    (feat_z, feat_h, feat_w), (valid_feat_z, valid_feat_h, valid_feat_w))
                multi_level_flags.append(flags)
            valid_flag_list.append(multi_level_flags)

        return anchor_list, valid_flag_list

    def loss_single(self, cls_score, bbox_pred, labels, label_weights,
                    bbox_targets, bbox_weights, anchors, num_total_samples, cfg, level, gt_bboxes, iteration):
        # classification loss
        labels = labels.reshape(-1)
        label_weights = label_weights.reshape(-1)
        # 3D images
        cls_score = cls_score.permute(0, 3, 4, 2, 1).reshape(-1, self.cls_out_channels)
        # debug only...
        # self.visualize_anchors_across_levels(anchors, gt_bboxes, labels, label_weights, iteration=iteration, level=level['level'])
        # self.print_cls_scores(cls_score, labels, label_weights, num_total_samples)

        if self.use_sigmoid_cls:
            if self.use_focal_loss:
                cls_criterion = weighted_sigmoid_focal_loss
            else:
                cls_criterion = weighted_binary_cross_entropy
        else:
            if self.use_focal_loss:
                raise NotImplementedError
            else:
                cls_criterion = weighted_cross_entropy
        if self.use_focal_loss:
            loss_cls = cls_criterion(
                cls_score,
                labels,
                label_weights,
                gamma=cfg.gamma,
                alpha=cfg.alpha,
                avg_factor=num_total_samples)
        else:
            loss_cls = cls_criterion(
                cls_score, labels, label_weights, avg_factor=num_total_samples)
        # regression loss
        bbox_targets = bbox_targets.reshape(-1, 6)
        bbox_weights = bbox_weights.reshape(-1, 6)
        # 3D images
        bbox_pred = bbox_pred.permute(0, 3, 4, 2, 1).reshape(-1, 6)

        loss_reg = weighted_smoothl1(
            bbox_pred,
            bbox_targets,
            bbox_weights,
            beta=cfg.smoothl1_beta,
            avg_factor=num_total_samples)

        # debug only...
        # print('level {} loss_cls: {}'.format(level['level'], loss_cls))
        # print('level {} loss_reg: {}'.format(level['level'], loss_reg))
        level['level'] += 1
        return loss_cls, loss_reg

    def loss(self,
             cls_scores,
             bbox_preds,
             gt_bboxes,
             gt_labels,
             img_metas,
             cfg,
             gt_bboxes_ignore=None,
             iteration=None):
        # 3D images
        featmap_sizes = [featmap.size()[-3:] for featmap in cls_scores]

        assert len(featmap_sizes) == len(self.anchor_generators)
        anchor_list, valid_flag_list = self.get_anchors(
            featmap_sizes, img_metas)
        sampling = False if self.use_focal_loss else True
        label_channels = self.cls_out_channels if self.use_sigmoid_cls else 1
        cls_reg_targets = anchor_target(
            anchor_list,
            valid_flag_list,
            gt_bboxes,
            img_metas,
            self.target_means,
            self.target_stds,
            cfg,
            gt_bboxes_ignore_list=gt_bboxes_ignore,
            gt_labels_list=gt_labels,
            label_channels=label_channels,
            sampling=sampling,
            iteration=iteration)
        if cls_reg_targets is None:
            return None
        (labels_list, label_weights_list, bbox_targets_list, bbox_weights_list,
         num_total_pos, num_total_neg, anchors_list, inside_flags) = cls_reg_targets
        self.pos_indices = inside_flags
        num_total_samples = (num_total_pos if self.use_focal_loss else
                             num_total_pos + num_total_neg)
        level = {'level': 0}
        losses_cls, losses_reg = multi_apply(
            self.loss_single,
            cls_scores,
            bbox_preds,
            labels_list,
            label_weights_list,
            bbox_targets_list,
            bbox_weights_list,
            anchors_list,
            num_total_samples=num_total_samples,
            cfg=cfg,
            level=level,
            gt_bboxes=gt_bboxes,
            iteration=iteration)
        return dict(loss_cls=losses_cls, loss_reg=losses_reg)

    def get_bboxes(self, cls_scores, bbox_preds, img_metas, cfg,
                   rescale=False, img_meta_2=None, img_meta_3=None):
        if img_meta_2 is not None:
            img_metas = img_meta_2
        if img_meta_3 is not None:
            img_metas = img_meta_3

        if self.pos_indices_test is None and hasattr(cfg, 'different_img_size') and cfg.different_img_size:
            featmap_sizes = [featmap.size()[-3:] for featmap in cls_scores]
            anchor_list, valid_flag_list = self.get_anchors(featmap_sizes, img_metas)
            inside_flags = get_anchor_inside_flags(anchor_list, valid_flag_list, img_metas)
            self.pos_indices_test = inside_flags

        assert len(cls_scores) == len(bbox_preds)
        num_levels = len(cls_scores)

        mlvl_anchors = [
            self.anchor_generators[i].grid_anchors(cls_scores[i].size()[-3:],
                                                   self.anchor_strides[i], self.anchor_strides_depth[i])
            for i in range(num_levels)
        ]
        result_list, anchors_list = [], []
        for img_id in range(len(img_metas)):
            cls_score_list = [
                cls_scores[i][img_id].detach() for i in range(num_levels)
            ]
            bbox_pred_list = [
                bbox_preds[i][img_id].detach() for i in range(num_levels)
            ]
            img_shape = img_metas[img_id]['img_shape']
            scale_factor = img_metas[img_id]['scale_factor']
            proposals, anchors = self.get_bboxes_single(cls_score_list, bbox_pred_list,
                                               mlvl_anchors, img_shape,
                                               scale_factor, cfg, rescale)
            result_list.append(proposals)
            anchors_list.append(anchors)
        return result_list, anchors_list

    '''
    for debugging 3D images only...
    '''
    def visualize_anchor_boxes(self, imgs, cls_scores, img_metas, slice_num=45, top_k=None, shuffle=False):
        featmap_sizes = [featmap.size()[-3:] for featmap in cls_scores]
        assert len(featmap_sizes) == len(self.anchor_generators)
        anchor_list, valid_flag_list = self.get_anchors(
            featmap_sizes, img_metas)
        img = tensor2img3D(imgs, slice_num=slice_num)
        anchors = []
        unique_set = set()
        for bboxes in anchor_list[0]:
            bboxes = bboxes.cpu().numpy()
            for bbox in bboxes:
                # select each aspect ratio bounding box in the middle of an image
                # if bbox[0] >= 100 and bbox[0] <= 400 and bbox[2] >= 100 and bbox[2] <= 400 and \
                #     bbox[1] >= 150 and bbox[1] <= 450 and bbox[3] >= 150 and bbox[3] <= 450 and \
                #     slice_num >= bbox[4] and slice_num <= bbox[5] and (bbox[5] - bbox[4]) not in unique_set:
                #     unique_set.add(bbox[5] - bbox[4])
                #     anchors.append([bbox[0], bbox[1], bbox[2], bbox[3]])

                # Get all anchors in the middle of the image
                if bbox[0] >= 100 and bbox[0] <= 400 and bbox[2] >= 100 and bbox[2] <= 400 and \
                    bbox[1] >= 150 and bbox[1] <= 450 and bbox[3] >= 150 and bbox[3] <= 450 and \
                    slice_num >= bbox[4] and slice_num <= bbox[5] and (bbox[2] - bbox[0]) not in unique_set:
                    anchors.append([bbox[0], bbox[1], bbox[2], bbox[3]])
        print(unique_set)
        breakpoint()
        if shuffle is True:
            while True:
                random.shuffle(anchors)
                mmcv.imshow_bboxes(img, np.array(anchors), top_k=20)
        elif top_k is None:
            mmcv.imshow_bboxes(img, np.array(anchors))
        else:
            mmcv.imshow_bboxes(img, np.array(anchors), top_k=top_k)
        breakpoint()

    def imwrite(self, img, file_path, params=None):
        """Write image to file
        """
        return cv2.imwrite(file_path, img, params)

    def visualize_anchors_across_levels(self,
                    bboxes,
                    gt_bboxes=None,
                    labels=None,
                    label_weights=None,
                    colors=[(0, 255, 0), (0, 0, 255), (255, 0, 0)],
                    top_k=-1,
                    thickness=1,
                    show=True,
                    iteration=None,
                    level=None):
        """Draw bboxes on an image.
        """
        if labels.dim() == 3:
            # pick the first patient for display
            cur_labels = labels[0, :, :]
            cur_label_weights = label_weights[0, :, :]
        else:
            cur_labels = labels
            cur_label_weights = label_weights

        pos_indices_labels = [index for index, val in enumerate(cur_labels.cpu().numpy()) if val > 0]
        labels_weights_indices = [index for index, val in enumerate(cur_label_weights.cpu().numpy()) if val > 0]

        # uncomment to print negative boxes
        neg_labels_indices = []
        for label_weight_index in labels_weights_indices:
            if label_weight_index not in pos_indices_labels:
                neg_labels_indices.append(label_weight_index)

        if bboxes.dim() == 3:
            pos_bboxes = bboxes[0, pos_indices_labels, :]
            neg_bboxes = bboxes[0, neg_labels_indices, :]
        else:
            pos_bboxes = bboxes[pos_indices_labels, :]
            neg_bboxes = bboxes[neg_labels_indices, :]

        pos_bboxes = pos_bboxes.cpu().numpy()
        pos_bboxes = [pos_bboxes]
        neg_bboxes = neg_bboxes.cpu().numpy()
        neg_bboxes = [neg_bboxes]
        gt_bboxes = gt_bboxes[0].cpu().numpy()
        gt_bboxes = [gt_bboxes]
        # keys = {'xy', 'xz', 'yz'}
        keys = {'xy'}

        for key in keys:
            if key == 'xy':
                img = np.zeros((768, 768, 3))
            elif key == 'xz':
                img = np.zeros((160, 768, 3))
            elif key == 'yz':
                img = np.zeros((160, 768, 3))

            if len(pos_bboxes) > 0:
                for i, _bboxes in enumerate(pos_bboxes):
                    _bboxes = _bboxes.astype(np.int32)
                    if top_k <= 0:
                        _top_k = _bboxes.shape[0]
                    else:
                        _top_k = min(top_k, _bboxes.shape[0])
                    for j in range(_top_k):
                        if key == 'xy':
                            left_top = (_bboxes[j, 0], _bboxes[j, 1])
                            right_bottom = (_bboxes[j, 2], _bboxes[j, 3])
                        elif key == 'xz':
                            left_top = (_bboxes[j, 0], _bboxes[j, 4])
                            right_bottom = (_bboxes[j, 2], _bboxes[j, 5])
                        elif key == 'yz':
                            left_top = (_bboxes[j, 1], _bboxes[j, 4])
                            right_bottom = (_bboxes[j, 3], _bboxes[j, 5])

                        cv2.rectangle(img, left_top, right_bottom, colors[0], thickness=thickness)

            if len(neg_bboxes) > 0:
                for i, _bboxes in enumerate(neg_bboxes):
                    _bboxes = _bboxes.astype(np.int32)
                    if top_k <= 0:
                        _top_k = _bboxes.shape[0]
                    else:
                        _top_k = min(top_k, _bboxes.shape[0])
                    for j in range(_top_k):
                        if key == 'xy':
                            left_top = (_bboxes[j, 0], _bboxes[j, 1])
                            right_bottom = (_bboxes[j, 2], _bboxes[j, 3])
                        elif key == 'xz':
                            left_top = (_bboxes[j, 0], _bboxes[j, 4])
                            right_bottom = (_bboxes[j, 2], _bboxes[j, 5])
                        elif key == 'yz':
                            left_top = (_bboxes[j, 1], _bboxes[j, 4])
                            right_bottom = (_bboxes[j, 3], _bboxes[j, 5])

                        cv2.rectangle(img, left_top, right_bottom, colors[2], thickness=thickness)

            if gt_bboxes is not None and len(gt_bboxes) > 0:
                # gt_bboxes
                for i, _gt_bboxes in enumerate(gt_bboxes):
                    _gt_bboxes = _gt_bboxes.astype(np.int32)
                    if top_k <= 0:
                        _top_k = _gt_bboxes.shape[0]
                    else:
                        _top_k = min(top_k, _gt_bboxes.shape[0])
                    for j in range(_top_k):
                        if key == 'xy':
                            left_top = (_gt_bboxes[j, 0], _gt_bboxes[j, 1])
                            right_bottom = (_gt_bboxes[j, 2], _gt_bboxes[j, 3])
                        elif key == 'xz':
                            left_top = (_gt_bboxes[j, 0], _gt_bboxes[j, 4])
                            right_bottom = (_gt_bboxes[j, 2], _gt_bboxes[j, 5])
                        elif key == 'yz':
                            left_top = (_gt_bboxes[j, 1], _gt_bboxes[j, 4])
                            right_bottom = (_gt_bboxes[j, 3], _gt_bboxes[j, 5])
                        cv2.rectangle( img, left_top, right_bottom, colors[1], thickness=1)
            self.imwrite(img, 'tests/iter_{}_lvl{}_{}.png'.format(iteration, level, key))

    def visualize_anchors_per_slice(self,
                    bboxes,
                    gt_bboxes=None,
                    bbox_weights=None,
                    colors=[(0, 255, 0), (0, 0, 255)],
                    top_k=-1,
                    thickness=1,
                    show=True,
                    iteration=None):
        """Draw bboxes on an image.
        """
        if bbox_weights.dim() == 3:
            # pick the first patient for display
            cur_bbox_weights = bbox_weights[0, :, :]
        else:
            cur_bbox_weights = bbox_weights

        pos_indices_bbox_weights = [index for index, bbox in enumerate(cur_bbox_weights.cpu().numpy()) if bbox[0] > 0]

        if bboxes.dim() == 3:
            bboxes = bboxes[0, pos_indices_bbox_weights, :]
        else:
            bboxes = bboxes[pos_indices_bbox_weights, :]

        bboxes = bboxes.cpu().numpy()
        bboxes = [bboxes]
        gt_bboxes = gt_bboxes[0].cpu().numpy()
        gt_bboxes = [gt_bboxes]
        each_slice_gt = [[] for i in range(160)]
        each_slice_pred = [[] for i in range(160)]

        if gt_bboxes is not None and len(gt_bboxes) > 0:
            # gt_bboxes
            for i, _gt_bboxes in enumerate(gt_bboxes):
                _gt_bboxes = _gt_bboxes.astype(np.int32)
                if top_k <= 0:
                    _top_k = _gt_bboxes.shape[0]
                else:
                    _top_k = min(top_k, _gt_bboxes.shape[0])
                for j in range(_top_k):
                    for cur_slice_num in range(_gt_bboxes[j, 4], _gt_bboxes[j, 5]+1):
                        # left top and right bottom
                        each_slice_gt[cur_slice_num].append(((_gt_bboxes[j, 0], _gt_bboxes[j, 1]), (_gt_bboxes[j, 2], _gt_bboxes[j, 3])))

        if len(bboxes) > 0:
            # anchors
            for i, _bboxes in enumerate(bboxes):
                _bboxes = _bboxes.astype(np.int32)
                if top_k <= 0:
                    _top_k = _bboxes.shape[0]
                else:
                    _top_k = min(top_k, _bboxes.shape[0])
                for j in range(_top_k):
                    for cur_slice_num in range(_bboxes[j, 4], _bboxes[j, 5]+1):
                        # left top and right bottom
                        each_slice_pred[cur_slice_num].append(((_bboxes[j, 0], _bboxes[j, 1]), (_bboxes[j, 2], _bboxes[j, 3])))
        slice_num = 0
        for slice_gt, slice_pred in zip(each_slice_gt, each_slice_pred):
            # we only print when gt_bbox exists in a slice
            if len(slice_gt) > 0:
                img = np.zeros((768, 768, 3))
                for coords in slice_gt:
                    # red
                    cv2.rectangle( img, coords[0], coords[1], colors[1], thickness=1)
                for coords in slice_pred:
                    # green
                    cv2.rectangle(img, coords[0], coords[1], colors[0], thickness=1)
                self.imwrite(img, 'tests2/iter_{}_slice_{}.png'.format(iteration, slice_num))
            slice_num += 1

    def _expand_binary_labels(self, labels, label_weights, label_channels):
        bin_labels = labels.new_full((labels.size(0), label_channels), 0)
        inds = torch.nonzero(labels >= 1).squeeze()
        if inds.numel() > 0:
            bin_labels[inds, labels[inds] - 1] = 1
        bin_label_weights = label_weights.view(-1, 1).expand(
            label_weights.size(0), label_channels)
        return bin_labels, bin_label_weights

    def print_cls_scores(self, cls_score, labels, label_weights, num_total_samples):
        pos_labels_indices = [index for index, val in enumerate(labels.cpu().numpy()) if val > 0]
        labels_weights_indices = [index for index, val in enumerate(label_weights.cpu().numpy()) if val > 0]

        # make sure every positive label index exist in positive label weights
        for pos_label_index in pos_labels_indices:
            if pos_label_index not in labels_weights_indices:
                breakpoint()

        neg_labels_indices = []
        for label_weight_index in labels_weights_indices:
            if label_weight_index not in pos_labels_indices:
                neg_labels_indices.append(label_weight_index)

        label, weight = self._expand_binary_labels(labels, label_weights, cls_score.size(-1))
        losses = F.binary_cross_entropy_with_logits(cls_score, label.float(), weight.float(), reduction='none')[None] / num_total_samples
        adjusted_scores = cls_score.sigmoid()
        pos_scores = adjusted_scores[pos_labels_indices]
        neg_scores = adjusted_scores[neg_labels_indices]
        out = open('output.json', 'a+')

        out.write("number of positive samples: {}\n".format(len(pos_labels_indices)))
        out.write("number of negative samples: {}\n".format(len(neg_labels_indices)))

        cur_counter = 0
        max_to_print = 20 # only print this number of entries
        for index, score, loss in zip(pos_labels_indices, pos_scores, losses.flatten()[pos_labels_indices]):
            if cur_counter == max_to_print:
                break
            out.write("positive sample's index {}, score {}, loss {}\n".format(index, score.data, loss))
            cur_counter += 1
        out.write("\n")

        cur_counter = 0
        for index, score, loss in zip(neg_labels_indices, neg_scores, losses.flatten()[neg_labels_indices]):
            if cur_counter == max_to_print:
                break
            out.write("negative sample's index {}, score {}, loss {}\n".format(index, score.data, loss))
            cur_counter += 1


from .losses import wind_mean_squared_error
from .util import agg_window,create_windowed_arr,save_multiple_graph,get_output,gather_auc_avg_per_tol,join_mean_std
import matplotlib.pyplot as plt
from keras.models import Sequential, Model
import numpy as np
import pandas as pd
import os

class Params(object):
    """
    Parameters class to handlde parameters for the ROI based models
    """
    def __init__(self,width=64, height=64,win_length=8,channels=1,dset='Thermal_track',d_type='frame',lambda_=1.0,regularizer_list=None,R_name=None,D_name=None,batch_size=32,break_win=10):
        self.width=width
        self.height=height
        self.win_length=win_length
        self.channels=channels
        self.dset=dset
        self.d_type=d_type
        self.batch_size = batch_size
        self.lambda_=lambda_
        self.regularizer_list=regularizer_list
        self.R_name=R_name
        self.D_name=D_name
        self.gap=break_win

    def create_model_name(self):
        return  self.get_model_type()+ '_{}'.format(str(self.lambda_))
    def create_hp_name(self):
        return  'lambda_{}'.format(str(self.lambda_))
    def get_model_type(self):
        R_name=self.R_name
        D_name=self.D_name

        return R_name+'_'+D_name
    def get_model_dir(self):
        return self.get_root_path()+'/models'
    def get_R_path(self,epochs_trained):
        return self.get_root_path()+'/models/GAN_R_weights_epoch-{}.h5'.format(epochs_trained)
    def get_D_path(self,epochs_trained):
        return self.get_root_path()+'/models/GAN_D_weights_epoch-{}.h5'.format(epochs_trained)
    def get_root_path(self):
        return './{}/{}/{}/{}'.format(self.dset,self.d_type,self.get_model_type(),self.create_hp_name())

class CAE_GAN3D(object):
    '''
        Class used to train and test the base adversarial model
    '''
    def __init__(self, train_par=None,stride=1):

        self.train_par=train_par
        self.stride=stride
        self.SHAPE = (self.train_par.width, self.train_par.height, train_par.channels)

    def initialize_model(self,Reconstructor , Discriminator ):
        print("Compiling GAN model.")
        self.R = Reconstructor
        self.D = Discriminator
        print('Discriminator')
        print(self.D.summary())

        print('Reconstructor')
        print(self.R.summary())
        self.OPTIMIZER = 'adam'

        self.stacked_R_D = self.stacked_R_D()
        loss_weights = {'D':1.0, 'decoded':self.train_par.lambda_}

        self.stacked_R_D.compile(loss={'D': 'binary_crossentropy', 'decoded': 'mean_squared_error'},\
         optimizer=self.OPTIMIZER, loss_weights = loss_weights)

    def stacked_R_D(self):
        '''
        Used for training Reconstructor. Dicscriminator is freezed.
        '''

        self.D.trainable = False
        model = Model(inputs = self.R.input, outputs = [self.R.output, self.D(self.R.output)],name='stacked')
        print('stacked')
        print(model.summary())

        return model

    def create_windowed_data(self, videos_dic,stride=1,data_key='FRAME'):
        '''
            Create windows of frames
        '''
        total = []
        img_width, img_height,channels,win_length=self.train_par.width,self.train_par.height,self.train_par.channels,self.train_par.win_length
        if data_key=='FLOW':
            win_length=win_length-1
        for vid_name in videos_dic.keys():
            # print('Video Name', vid_name)
            vid_windowed_list=[]
            sub_vid_list=videos_dic[vid_name][data_key]
            for sub_vid in sub_vid_list:
                vid_windowed_list.append(create_windowed_arr(sub_vid, stride, win_length))
            # print("Number of sub videos: ",len(vid_windowed_list))
            vid_windowed=np.concatenate(vid_windowed_list)

            total.append(vid_windowed)
        total=np.concatenate(total)
        # print("Windowed data shape:")
        # print(total.shape)
        return total

    def get_MSE_all_agg(self, test_data,type=['frame','window']):
        """
            Anomaly scores based on MSE
        """

        img_width, img_height, win_length, channels,model = self.train_par.width, self.train_par.height ,self.train_par.win_length, self.train_par.channels, self.R
        if self.train_par.d_type=='opticalFLow':
            win_length=win_length-1
        recons_seq = model.predict([test_data]) #(samples-win_length+1, win_length, wd,ht,1)
        print(recons_seq.shape)
        RE=wind_mean_squared_error(test_data,recons_seq,win_length, img_height,img_width,channels)
        print('RE.shape', RE.shape)

        RE_dict = {}
        agg_type_list=[]
        if 'frame' in type:
            agg_type_list.append('x_std')
            agg_type_list.append('x_mean')
        if 'window' in type:
            agg_type_list.append('in_std')
            agg_type_list.append('in_mean')

        for agg_type in agg_type_list:

            RE_dict[agg_type] = agg_window(RE, agg_type)

        return RE_dict, recons_seq

    def train(self, X_train_frame,epochs= 500,epochs_trained=0, save_interval = 10):
        '''
            Train the adversarial framework
            X_train_frame- window of frames
        '''
        print('Using save root:', self.train_par.get_root_path())
        self.save_root=self.train_par.get_root_path()
        batch_size=self.train_par.batch_size
        print('self.stacked_R_D.metrics_names', self.stacked_R_D.metrics_names)
        print('self.D.metrics_names', self.D.metrics_names)
        num_batches = int(X_train_frame.shape[0]/batch_size)
        print("Train frame dataset shape",X_train_frame.shape)
        print("Number of batches",num_batches)
        #model save dir
        if not os.path.isdir(self.train_par.get_model_dir()):
            os.makedirs(self.train_par.get_model_dir())
        d_loss_list = [] # Discriminator loss
        r_loss_list_RE = [] #Reconstruction error
        r_loss_list_BCE = [] #Binary cross entropy
        loss_root = self.save_root + '/loss'

        #Creating loss directory
        if not os.path.isdir(loss_root):
            print("Creating loss directory ")
            os.makedirs(loss_root)
        print("Loss file status................")
        if os.path.isfile(loss_root + '/d_loss_epoch-{}.npy'.format(epochs_trained)):
            print("D Loss file found")
            d_loss_list=list(np.load(loss_root + '/d_loss_epoch-{}.npy'.format(epochs_trained)))
        if os.path.isfile(loss_root + '/r_loss_RE_epoch-{}.npy'.format(epochs_trained)):
            print("RE Loss file found")
            r_loss_list_RE=list(np.load(loss_root + '/r_loss_RE_epoch-{}.npy'.format(epochs_trained)))
        if os.path.isfile(loss_root + '/r_loss_BCE_epoch-{}.npy'.format(epochs_trained)):
            print("BCE Loss file found")
            r_loss_list_BCE=list(np.load(loss_root + '/r_loss_BCE_epoch-{}.npy'.format(epochs_trained)))

        for epoch in range(epochs_trained+1,epochs):
            ## train discriminator

            random_index =  np.random.randint(0, len(X_train_frame) - batch_size)

            permutated_indexes = np.random.permutation(X_train_frame.shape[0])

            for step in range(num_batches):

                batch_indeces = permutated_indexes[step*batch_size:(step+1)*batch_size]

                legit_images = X_train_frame[batch_indeces]
                #R Input
                recons_images = self.R.predict([legit_images])

                x_combined_batch_size = np.concatenate((legit_images,recons_images))
                y_combined_batch_size = np.concatenate((np.ones((batch_size, 1)), np.zeros((batch_size, 1))))

                # First train Discriminator
                d_loss = self.D.train_on_batch(x_combined_batch_size, y_combined_batch_size)

                d_loss_list.append(d_loss)

                # Train Reconstructor
                y_mislabled = np.ones((batch_size, 1))
                r_loss = self.stacked_R_D.train_on_batch([legit_images], {'decoded':legit_images,'D':y_mislabled})

                r_loss_list_RE.append(r_loss[1])
                r_loss_list_BCE.append(r_loss[2])

                if step % 10 == 0:

                    print('epoch: {}, step {}, [Discriminator :: d_loss: {}], [ Reconstructor :: RE loss, BCE loss: {}, {}]'.format(epoch, step, d_loss, r_loss[1], r_loss[2]))

            if epoch % save_interval == 0 or epoch == epochs-1:

                save_string = self.train_par.get_R_path(epoch)
                self.R.save_weights(save_string)

                save_string = self.train_par.get_D_path(epoch)
                self.D.save_weights(save_string)
                print('saving images')
                np.random.seed(0)
                test_idxs = np.random.choice(len(X_train_frame), 8, replace = False)

                test_ims = X_train_frame[test_idxs]
                print(test_ims.shape)
                if self.train_par.d_type=='frame':
                    self.plot_images_3D(save2file=True, step=epoch, test_window = test_ims,d_type='thermal')
                elif self.train_par.d_type=='opticalFLow':
                    self.plot_images_3D(save2file=True, step=epoch, test_window = test_ims,d_type='flow')
                else:
                    print("Invalid data type in Params")

                np.save(loss_root + '/d_loss_epoch-{}.npy'.format(epoch), np.array(d_loss_list))
                np.save(loss_root +  '/r_loss_RE_epoch-{}.npy'.format(epoch), np.array(r_loss_list_RE))
                np.save(loss_root + '/r_loss_BCE_epoch-{}.npy'.format(epoch), np.array(r_loss_list_BCE))
                save_multiple_graph(x_list=[r_loss_list_RE,r_loss_list_BCE,d_loss_list],labels=['R_RE','R_BCE','D_loss'],x_label='Batches',y_label='Losses',title='Loss Plot',path=loss_root+'/log_loss.png',log_plot=True)
                save_multiple_graph(x_list=[d_loss_list,r_loss_list_RE,r_loss_list_BCE],labels=['D_loss','R_RE','R_BCE'],x_label='Batches',y_label='Losses',title='Loss Plot',path=loss_root+'/loss.png',log_plot=False)

    def plot_images_3D(self, save2file=False,  samples=16, step=0, test_window = None,d_type=None):
        '''
            Visualization of input and reconstrcuted sequence. Save or save 4th frame of the input and output windows.
        '''
        test_ims = test_window[:,4,:,:,:]
        img_root = self.save_root + "/"+d_type+"_images/"
        channels=self.train_par.channels
        rec_images = self.R.predict([test_window])

        if not os.path.isdir(img_root):
            os.makedirs(img_root)

        filename = img_root + "/img_{}.png".format(step)

        rec_images = rec_images[:,4,:,:,:]

        if d_type=='flow' and channels==3:
            rec_images=rec_images[:,:,:,2:3]
            test_ims=test_ims[:,:,:,2:3]
            channels=1

        plt.figure(figsize=(10,10))
        for i in range(rec_images.shape[0]):
            plt.subplot(4, 4, i+1)
            image = rec_images[i, :, :, :]
            if channels==3:
                image = np.reshape(image, [ self.train_par.height, self.train_par.width, channels])
            else:
                image = np.reshape(image, [ self.train_par.height, self.train_par.width])
            plt.imshow(image, cmap='gray')
            plt.axis('off')
            plt.tight_layout()

        for i in range(test_ims.shape[0]):
            i+=8
            plt.subplot(4, 4, i+1)
            image = test_ims[i-8, :, :, :]
            if channels==3:
                image = np.reshape(image, [ self.train_par.height, self.train_par.width, channels])
            else:
                image = np.reshape(image, [ self.train_par.height, self.train_par.width])
            plt.imshow(image, cmap='gray')
            plt.axis('off')
            plt.tight_layout()

        if save2file:
            plt.savefig(filename)
            plt.close('all')
        else:
            plt.show()

    def test(self, test_videos, score_type = 'R', epochs = None,plot=False,tolerance_limit=8):

        '''
        Gets AUC ROC/PR for all videos using RE.

        choose score type 'R'
        '''

        dset, img_width, img_height, win_length  = self.train_par.dset,  self.train_par.width, self.train_par.height, self.train_par.win_length
        self.save_root = self.train_par.get_root_path() + '/testing/epochs_{}'.format(epochs)
        stride = self.stride

        model_name = self.train_par.create_model_name()
        model_name += '_{}'.format(score_type)

        print(model_name)

        aucs = []
        std_total = []
        mean_total = []
        labels_total_l = []
        i = 0 #vid index TODO rename
        num_vids = len(test_videos)
        print('num of test vids', num_vids)

        ROC_mat = np.zeros((num_vids,2*tolerance_limit+2)) # 35 is num_vids, 20 scores-Xstd,Xmean,tols std..,tols mean..
        PR_mat = np.zeros((num_vids,2*tolerance_limit+2))
        print('score_type', score_type)

        for vid_name in  test_videos.keys():
            print("--------------------------")
            print("--------------------------")
            print("Processing ",vid_name)
            print("--------------------------")
            print("--------------------------")

            # vid_total, labels_total = restore_Fall_vid(data_dict, Fall_name, NFF_name)
            vid_total_list=None
            frame_numbers_list=None
            start,end=None,None
            labels_total_list=test_videos[vid_name]['LABELS']

            display_name = vid_name

            test_labels = np.concatenate(labels_total_list)

            print("Number of Labels",len(test_labels))

            if self.train_par.d_type=='frame':
                vid_total_list=test_videos[vid_name]['FRAME']
                frame_numbers_list=test_videos[vid_name]['NUMBER']
                frame_numbers = np.concatenate(frame_numbers_list)
                print("Number of frames",len(frame_numbers))
                start,end=test_videos[vid_name]['START_END']

                test_data_list = [vid.reshape(len(vid), img_width, img_height, self.train_par.channels) for vid in vid_total_list]
                test_data_windowed_list = [create_windowed_arr(test_data, stride, win_length) for test_data in test_data_list]#create_windowed_arr function in data_management.py

            elif self.train_par.d_type=='opticalFLow':
                vid_total_list=test_videos[vid_name]['FLOW']
                #Create windows of length win_length-1
                test_data_list = [vid.reshape(len(vid), img_width, img_height, self.train_par.channels) for vid in vid_total_list]
                test_data_windowed_list = [create_windowed_arr(test_data, stride, win_length-1) for test_data in test_data_list]#create_windowed_arr function in data_management.py

            else:
                print("Invlaid d_type in the params")
            num_sub_videos=len(test_data_windowed_list)
            print("Number of sub flow videos,",num_sub_videos)

            if score_type == 'R':
                in_mean_RE=[]
                in_std_RE=[]
                x_std_RE=[]
                x_mean_RE=[]
                for index in range(num_sub_videos):
                    test_data_windowed=test_data_windowed_list[index]

                    RE_dict, recons_seq = self.get_MSE_all_agg(test_data_windowed) #Return dict with value for each score style
                    in_mean_RE.append(RE_dict['in_mean'])
                    in_std_RE.append(RE_dict['in_std'])

                    x_std_RE.append(RE_dict['x_std'])
                    x_mean_RE.append(RE_dict['x_mean'])

                in_mean_RE=np.concatenate(in_mean_RE)
                in_std_RE=np.concatenate(in_std_RE)
                x_std_RE=np.concatenate(x_std_RE)
                x_mean_RE=np.concatenate(x_mean_RE)

                final_in_mean = in_mean_RE
                final_in_std = in_std_RE

            #frame based scores only for thermal frames
            if self.train_par.d_type=='frame':
                auc_x_std, conf_mat, g_mean, ap_x_std = get_output(labels = test_labels,\
                    predictions = x_std_RE, data_option = 'NA', to_plot = False)
                auc_x_mean, conf_mat, g_mean, ap_x_mean = get_output(labels = test_labels,\
                    predictions = x_mean_RE, data_option = 'NA', to_plot = False)

                ROC_mat[i,0] = auc_x_std
                ROC_mat[i,1] = auc_x_mean

                PR_mat[i,0] = ap_x_std
                PR_mat[i,1] = ap_x_mean

            #window based scores
            # print('final_in_mean.shape', final_in_mean.shape, 'final_in_std.shape', final_in_std.shape)
            tol_mat, tol_keys = gather_auc_avg_per_tol(final_in_mean, final_in_std, labels_list = labels_total_list, win_len = win_length,tolerance_limit=tolerance_limit)
            AUROC_tol = tol_mat[0]
            AUPR_tol = tol_mat[1]
            num_scores_tol = tol_mat.shape[1]

            for k in range(num_scores_tol):
                j = k+2 #start at 2, first two were for X_std and X_mean
                ROC_mat[i,j] = AUROC_tol[k]
                PR_mat[i,j] = AUPR_tol[k]

            i += 1

            if plot == True and score_type == 'R' and self.train_par.d_type=='frame':
                plt.plot(frame_numbers,x_std_RE, label='RE_std',linestyle='--', marker='.')
                plt.plot(frame_numbers,x_mean_RE, label='RE_mean',linestyle='--', marker='.')
                # plt.xticks([i+1 for i in range(max(frame_numbers))])
                plt.xlim(1,max(frame_numbers))
                # plt.ylim(0,1)
                plt.legend()
                plt.axvspan(start,end, alpha = 0.5)

                plot_save_p = self.save_root + '/scores_plots/'
                if not os.path.isdir(plot_save_p):
                    os.makedirs(plot_save_p)
                plt.savefig(plot_save_p + '{}.jpg'.format(vid_name))
                plt.close()

        #    break
        AUROC_avg = np.mean(ROC_mat, axis = 0)
        AUROC_std = np.std(ROC_mat, axis = 0)
        AUROC_avg_std = join_mean_std(AUROC_avg, AUROC_std)
        # print(AUROC_std)
        AUPR_avg = np.mean(PR_mat, axis = 0)
        AUPR_std = np.std(PR_mat, axis = 0)

        AUPR_avg_std = join_mean_std(AUPR_avg, AUPR_std)
        total = np.vstack((AUROC_avg_std, AUPR_avg_std))
        total_no_std = np.vstack((AUROC_avg, AUPR_avg))

        df = pd.DataFrame(data = total, index = ['AUROC','AUPR'], columns = ['X-STD','X-Mean'] + tol_keys)
        df_no_std = pd.DataFrame(data = total_no_std, index = ['AUROC','AUPR'], columns = ['X-STD','X-Mean'] + tol_keys)

        print(df)
        print(df_no_std)

        if not os.path.isdir(self.save_root):
            os.makedirs(self.save_root)

        save_path = self.save_root + '/AUC_{}.csv'.format(score_type)
        save_path_no_std = self.save_root + '/AUC_{}_no_std.csv'.format(score_type)

        print(save_path)
        df.to_csv(save_path)
        df_no_std.to_csv(save_path_no_std)


# -*- coding: utf-8 -*-

import os
import sys
import time
import math
import argparse
import numpy as np

import torch
import torch.nn as nn
from torch import cuda
import torch.nn.functional as F
from transformers import MBart50Tokenizer
from transformers import MBartForConditionalGeneration

sys.path.append("")
from utils.dataset import collate_fn
from utils.dataset import read_insts
from utils.optim import ScheduledOptim

filter_sizes = [1, 2, 3, 4, 5]
num_filters = [128, 128, 128, 128, 128]

os.environ["TOKENIZERS_PARALLELISM"] = "false"
device = 'cuda' if cuda.is_available() else 'cpu'

def load_embedding(tokenizer, embed_dim, embed_path=None):
    '''Parse an embedding text file into an torch.nn.Embedding layer.'''

    embedding = np.random.normal(scale=embed_dim ** -0.5,
                                 size=(len(tokenizer), embed_dim))
    if embed_path == None:
        return embedding

    print('[Info] Loading embedding')
    embed_dict = {}
    with open(embed_path) as file:
        for i, line in enumerate(file):
            if i == 0:
                continue
            tokens = line.rstrip().split()
            try:
                embed_dict[tokens[0]] = np.asarray(tokens[1:], dtype='float32')
            except:
                continue

    for i in range(len(tokenizer)):
        try:
            word = tokenizer.decode(i)
            if word in embed_dict:
                embedding[i] = embed_dict[word]
        except:
            print(i)

    return embedding

def evaluate(model, valid_loader, loss_fn, epoch):
    '''Evaluation function for style classifier'''
    model.eval()
    total_acc = 0.
    total_num = 0.
    total_loss = []
    with torch.no_grad():
        for batch in valid_loader:
            x_batch, y_batch = map(lambda x: x.to(device), batch)
            logits = model(x_batch)
            total_loss.append(loss_fn(logits, y_batch).item())
            _, y_hat = torch.max(logits,dim=-1)
            same = [float(p == q) for p, q in zip(y_batch, y_hat)]
            total_acc += sum(same)
            total_num += len(y_batch)
    model.train()
    print('[Info] Epoch {:02d}-valid | '
          'acc {:.2f}% | loss {:.4f}'.format(epoch,
          total_acc/total_num*100, np.mean(total_loss)))

    return total_acc/total_num, np.mean(total_loss)

class SCDataset(torch.utils.data.Dataset):
    def __init__(self, insts, label):
        self.insts = insts
        self.label = label

    def __getitem__(self, index):
        return self.insts[index], self.label[index]

    def __len__(self):
        return len(self.insts)

def SCIterator(insts_neg, insts_pos, pad_id, shuffle, opt):
    '''Data iterator for style classifier'''

    def cls_fn(insts):
        src, tgt = list(zip(*insts))
        src = collate_fn(src, pad_id)
        tgt = torch.LongTensor(tgt)

        return (src, tgt)

    num = len(insts_neg) + len(insts_pos)
    loader = torch.utils.data.DataLoader(
        SCDataset(
            insts=insts_neg + insts_pos,
            label=[0 if i < len(insts_neg)
                   else 1 for i in range(num)]),
        num_workers=2,
        shuffle=shuffle,
        collate_fn=cls_fn,
        batch_size=opt.batch_size)

    return loader

class EmbeddingLayer(nn.Module):
    def __init__(self, vocab_size, embed_dim, embeding):
        super(EmbeddingLayer, self).__init__()
        self.embeding = nn.Embedding(vocab_size, embed_dim)
        if embeding is not None:
            self.embeding.weight.data = embeding

    def forward(self, x):
        if len(x.size()) == 2:
            y = self.embeding(x)
        else:
            y = torch.matmul(x, self.embeding.weight)
        return y

class TextCNN(nn.Module):
    '''A TextCNN Classification Model'''

    def __init__(self, embed_dim, vocab_size, filter_sizes,
                 num_filters, embedding=None, dropout=0.0):
        super(TextCNN, self).__init__()

        self.feature_dim = sum(num_filters)
        self.embeder = EmbeddingLayer(vocab_size, embed_dim, embedding)
        self.convs = nn.ModuleList([
            nn.Conv2d(1, n, (f, embed_dim))
            for (n, f) in zip(num_filters, filter_sizes)
        ])

        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Sequential(
            self.dropout,
            nn.Linear(self.feature_dim, int(self.feature_dim / 2)), nn.ReLU(),
            nn.Linear(int(self.feature_dim / 2), 2)
        )

    def forward(self, inp):
        inp = self.embeder(inp).unsqueeze(1)
        convs = [F.relu(conv(inp)).squeeze(3) for conv in self.convs]
        pools = [F.max_pool1d(conv, conv.size(2)).squeeze(2) for conv in convs]
        out = torch.cat(pools, 1)
        logit = self.fc(out)

        return logit

    def build_embeder(self, vocab_size, embed_dim, embedding=None):
        embeder = nn.Embedding(vocab_size, embed_dim)
        nn.init.normal_(embeder.weight, mean=0, std=embed_dim ** -0.5)
        if embedding is not None:
            embeder.weight.data = torch.FloatTensor(embedding)

        return embeder

def main():
    parser = argparse.ArgumentParser('TextCNN Style Classifier')
    parser.add_argument('-lr', default=1e-3, type=float, help='learning rate')
    parser.add_argument('-dim', default=300, type=int, help='the embedding size')
    parser.add_argument('-style', default=0, type=int, help='transfer inf. to for.')
    parser.add_argument('-dataset', default='xformal', type=str, help='dataset name')
    parser.add_argument('-seed', default=42, type=int, help='pseudo random number seed')
    parser.add_argument('-lang', default='it_IT', type=str, help='the name of language')
    parser.add_argument('-dropout', default=.5, type=float, help="Keep prob in dropout.")
    parser.add_argument('-log_step', default=100, type=int, help='print log every x steps')
    parser.add_argument('-epoch', default=10, type=int, help='force stop at specified epoch')
    parser.add_argument('-eval_step', default=1000, type=int, help='early stopping training')
    parser.add_argument('-batch_size', default=24, type=int, help='maximum sents in a batch')

    opt = parser.parse_args()
    torch.manual_seed(opt.seed)
    print('[Info]', opt)

    model_path = "facebook/mbart-large-50"
    tokenizer = MBart50Tokenizer.from_pretrained(model_path, src_lang=opt.lang)
    pad_id = tokenizer.pad_token_id

    #if os.path.exists('checkpoints/embedding.pt'):
    #    embedding = torch.load('checkpoints/embedding.pt')
    #else:
    #    embed_path = '../../data/glove.840B.300d.txt'
    #    embedding = load_embedding(tokenizer, 300, embed_path)
    #    torch.save(embedding, 'checkpoints/embedding.pt')

    model = TextCNN(opt.dim, len(tokenizer), filter_sizes,
                    num_filters, None, opt.dropout)
    model.to(device).train()
    print('[Info] Built a model with {} parameters'.format(
           sum(p.numel() for p in model.parameters())))

    train_src, train_tgt, valid_src, valid_tgt = [], [], [], []
    for lang in [opt.lang]:
        train_src_lang, train_tgt_lang = read_insts(opt.dataset, opt.style, 'train', lang)
        valid_src_lang, valid_tgt_lang = read_insts(opt.dataset, opt.style, 'valid', lang)
        train_src = train_src + train_src_lang
        train_tgt = train_tgt + train_tgt_lang
        valid_src = valid_src + valid_src_lang
        valid_tgt = valid_tgt + valid_tgt_lang
    print('[Info] {} instances from train set'.format(len(train_src)))
    print('[Info] {} instances from valid set'.format(len(valid_tgt)))
    train_loader = SCIterator(train_src, train_tgt, pad_id, True, opt)
    valid_loader = SCIterator(valid_src, valid_tgt, pad_id, False,opt)

    loss_fn = nn.CrossEntropyLoss()
    optimizer = ScheduledOptim(
        torch.optim.Adam(filter(lambda x: x.requires_grad, model.parameters()),
                         betas=(0.9, 0.98), eps=1e-09), opt.lr, 1000)

    tab = 0
    avg_acc = 0
    corre_num = 0.
    total_num = 0.
    loss_list = []
    start = time.time()
    for epoch in range(opt.epoch):

        for batch in train_loader:
            src, tgt = map(lambda x: x.to(device), batch)
            optimizer.zero_grad()
            outs = model(src)
            loss = loss_fn(outs, tgt)
            loss.backward()
            optimizer.step()

            y_hat = outs.argmax(dim=-1)
            same = [int(p == q) for p, q in zip(tgt, y_hat)]
            corre_num += sum(same)
            total_num += len(tgt)
            loss_list.append(loss.item())

            if optimizer.cur_step % opt.log_step == 0:
                lr = optimizer._optimizer.param_groups[0]['lr']
                print('[Info] Epoch {:02d}-{:05d} | acc {:.2f}% | '
                    'loss {:.4f} | lr {:.5f} | second {:.2f}'.format(
                    epoch, optimizer.cur_step, corre_num / total_num * 100,
                    np.mean(loss_list), lr, time.time() - start))
                corre_num = 0.
                total_num = 0.
                loss_list = []
                start = time.time()

            if optimizer.cur_step % opt.eval_step == 0:
                valid_acc, valid_loss = evaluate(model, valid_loader, loss_fn, epoch)
                if avg_acc < valid_acc:
                    avg_acc = valid_acc
                    save_path = 'checkpoints/textcnn_{}.chkpt'.format(opt.lang)
                    torch.save(model.state_dict(), save_path)
                    print('[Info] The checkpoint file has been updated.')
                    tab = 0
                else:
                    tab += 1
                    if tab == 5:
                        exit()

if __name__ == '__main__':
    main()


#!/usr/bin/python3

import unittest
import inspect
#import matplotlib.pyplot as plt
from selenium import webdriver
import time
from random import randint
import numpy as np
#import tensorflow as tf
from scipy.stats import kurtosis, skew, gengamma

URL_VAALIKONE        = "https://www.vaalikone.fi/eduskunta2019/"
OUTPUT_FILE          = "/tmp/LOG"

_SOME_COMMON_DIV_    = "//*[@id='app']/div/main/div/div"
XPATH_QUESTION       = _SOME_COMMON_DIV_ + "[1]/div/div/div[1]/div[1]/div"
XPATH_ANSWER         = _SOME_COMMON_DIV_ + "[1]/div/div/div[1]/div[2]/div[1]/button[%d]"
XPATH_RESULTS_BUTTON = _SOME_COMMON_DIV_ + "[1]/div/div[2]/button"
XPATH_MATCH_PARTY    = _SOME_COMMON_DIV_ + "[1]/div/div[1]/div[3]/div[1]/div/div[1]/div/div[2]/div[2]"
XPATH_DISTRICT       = "//button[@data-district-name='Helsinki']"
CLASS_START_BUTTON   =  "big__1kEOQ"
CLASS_RESULTS_BUTTON =  "button__1ARU2"

class Test(unittest.TestCase):

    driver = None
    out = None

    def _init_driver(self):
        print(inspect.stack()[0][3])
        options = webdriver.ChromeOptions()
        options.add_argument('--ignore-certificate-errors')
        options.add_argument("--test-type")
        options.add_argument('--window-size=1420,1080')
        #options.add_argument('--no-sandbox')
        #options.add_argument('--headless')
        options.add_argument('--disable-gpu')
        options.add_argument('--disable-extensions')
        #options.binary_location = "/snap/bin/chromium"
        options.binary_location = "./chromedriver"
        self.driver = webdriver.Chrome(options=options)
        print(inspect.stack()[0][3] + " - OK")

    def _init_output(self):
        print(inspect.stack()[0][3])
        self.out = open(OUTPUT_FILE, "w")

    def _get_answer_set(self):
        print(inspect.stack()[0][3])
        x = np.random.gamma(2, 2, 30)
        print(x)
        #ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf')
        #x = np.random.normal(0, 2, 10000)   # create random values based on a normal distribution
        #print( 'excess kurtosis of normal distribution (should be 0): {}'.format( kurtosis(x) ))
        #print( 'skewness of normal distribution (should be 0): {}'.format( skew(x) ))
        return x

    def _select_district(self):
        print(inspect.stack()[0][3])
        time.sleep(2)
        el = self.driver.find_element_by_xpath(XPATH_DISTRICT)
        self.output("D", el.text)
        el.click()

    def _goto_main_page(self):
        print(inspect.stack()[0][3])
        self.driver.get(URL_VAALIKONE)

    def _goto_questions(self):
        print(inspect.stack()[0][3])
        time.sleep(2)
        self.driver.find_element_by_class_name(CLASS_START_BUTTON).click()

    def _loop_through_questions(self):
        print(inspect.stack()[0][3])
        time.sleep(2)
        for i in [randint(1, 5) for i in range(0, 30)]:
            time.sleep(1)
            self.output("Q", str(i) + " " + self.driver.find_element_by_xpath(XPATH_QUESTION).text)
            self.driver.find_element_by_xpath(XPATH_ANSWER % i).click()

    def _ok_startup_alert(self):
        print(inspect.stack()[0][3])
        time.sleep(5)
        self.driver.switchTo().alert().accept();

    def output(self, label, text = '', delim = ' '):
        line = str(delim.join([label, text]).encode('utf-8').strip())
        self.out.write(line + "\n")
        print(line)

    def _print_results(self):
        print(inspect.stack()[0][3])
        time.sleep(1)
        el = self.driver.find_element_by_xpath(XPATH_MATCH_PARTY)
        self.output("R", el.text.split("\n")[0])

    def _goto_results(self):
        print(inspect.stack()[0][3])
        try:
            time.sleep(1)
            el = self.driver.find_element_by_class_name(CLASS_RESULTS_BUTTON)
            time.sleep(1)
            el.click()
        except:
            self.output('Error in _goto_results')

    def _accept_cookies(self):
        print(inspect.stack()[0][3])
        time.sleep(2)
        self.driver.find_element_by_xpath("//*[@id='sccm-opt-out-c1']").click()
    def setUp(self):
        print(inspect.stack()[0][3])
        self._init_driver()
        self._init_output()

    def testOne(self):
        print(inspect.stack()[0][3])
        if True:
            self._get_answer_set()
        else:
            #self._ok_startup_alert()
            self._goto_main_page()
            self._accept_cookies()
            self._select_district()
            self._goto_questions()
            self._loop_through_questions()
            self._goto_results()
            self._print_results()

    def tearDown(self):
        print(inspect.stack()[0][3])
        if self.driver:
            self.driver.close()
        if self.out:
            self.out.close()

if __name__ == '__main__':
    unittest.main()


# -*- coding: utf-8 -*-#
from data_source import DataSource
import os
from utils import retry, s3_download
from urllib.request import urlopen
from io import BytesIO
from scipy.misc.pilutil import imread, imsave
import tensorflow as tf
import numpy as np
from PIL import Image
import requests

# vgg16 performs better in predicting nightlights but produces worse scoring features
MODEL = 'google_simple.h5'  # google_vgg16.h5 (much slower)
# TODO: we can allow a parameter in the config for the model to use, as long as the layer is called 'features'
LAYER = 'features'  # features
IMG_SIZE = 256
# Google images parameters.
ZOOM_LEVEL = 18
RAW_SIZE = 380
map_size = "500x500"
imagery_set = "satellite"

class GoogleImages(DataSource):
    """overloading the DataSource class."""

    def __init__(self, directory):
        DataSource.__init__(self, directory)

        """ Overload the directory path. """
        self.directory = os.path.join(self.directory, 'Google/')
        if not os.path.exists(self.directory):
            os.makedirs(self.directory)

        """ loads the model. """
        if os.path.exists('../Models/{}'.format(MODEL)) is False:
            print("INFO: downloading model. ")
            with requests.get('https://hrm-models.s3.eu-central-1.amazonaws.com/{}'.format(MODEL), stream=True) as r:
                with open('../Models/{}'.format(MODEL), 'wb') as f:
                    f.write(r.content)

        print("INFO: loading model for Google Images ...")
        self.net = tf.keras.models.load_model('../Models/{}'.format(MODEL), compile=False)
        self.net = tf.keras.models.Model(
            inputs=self.net.input,
            outputs=self.net.get_layer(LAYER).output)

    def download(self, lon, lat, step=False):
        """ given the list of coordinates, it downloads the corresponding satellite images.

        Notes:
            Google Maps Static API takes {latitude,longitude}

        Args:
            lon (list): list of longitudes.
            lat (list): list of latitudes.
            step (bool): if you want to add buffer images. SMore accurate but slow.
        """
        if step:
            print('INFO: adding steps to coordinates set.')
            lon, lat = self.add_steps(lon, lat)

        @retry(Exception, tries=4)
        def _urlopen_with_retry(url):
            return urlopen(url).read()

        _cnt, _total = 0, len(lon)  # counter and total number of images.

        for i, j in zip(lon, lat):

            print("INFO: {} images downloaded out of {}".format(_cnt, _total), end='\r')
            _cnt += 1

            file_name = str(i) + '_' + str(j) + '_' + str(ZOOM_LEVEL) + '.jpg'

            if os.path.exists(self.directory + file_name):
                print("INFO: {} already downloaded".format(file_name), end='\r')
            else:
                center_point = str(j) + "," + str(i)

                url = """https://maps.googleapis.com/maps/api/staticmap?center={}&zoom={}&size={}&maptype={}&key={}""".\
                    format(center_point, ZOOM_LEVEL, map_size, imagery_set, os.environ['Google_key'])

                buffer = BytesIO(_urlopen_with_retry(url))

                image = imread(buffer, mode='RGB')
                if (image[:, :, 0] == 245).sum() >= 100000:  # Gray image in Bing
                    print("No image in Bing API", file_name)
                else:
                    print('file path: ', os.path.join(self.directory, file_name))
                    imsave(os.path.join(self.directory, file_name), image[50:450, :, :])

    def featurize(self, lon, lat, step=0):
        """ Given lon lat lists, it extract the features from the image (if there) using the NN.

        Args:
            lon (list): list of longitudes.
            lat (list): list of latitudes.
            step (bool): if you want to add buffer images (9 in total). More accurate but slow.

        Returns:
            covariates for the coordinates pair.
        """
        if step:
            print('INFO: adding steps to coordinates set.')
            lon, lat = self.add_steps(lon, lat)

        _cnt, _total = 0, len(lon)  # counter and total number of images.

        features = []
        for i, j in zip(lon, lat):
            _cnt += 1

            file_name = str(i) + '_' + str(j) + '_' + str(ZOOM_LEVEL) + '.jpg'
            img_path = os.path.join(self.directory, file_name)

            image = Image.open(img_path, 'r')
            image = image.crop((  # crop center
                int(image.size[0] / 2 - RAW_SIZE / 2),
                int(image.size[1] / 2 - RAW_SIZE / 2),
                int(image.size[0] / 2 + RAW_SIZE / 2),
                int(image.size[1] / 2 + RAW_SIZE / 2)
            ))
            image = image.resize((IMG_SIZE, IMG_SIZE), Image.ANTIALIAS)
            image = np.array(image)/ 255.
            features.append(self.net.predict(np.array(image).reshape(1, IMG_SIZE, IMG_SIZE, 3)))

            if _cnt % 10 == 0: print("Feature extraction : {} tiles out of {}".format(_cnt, _total), end='\r')

        if step:
            # take the average for the 9 images around the original lon,lat
            f = np.copy(features)
            features = []
            for c in range(0, int(len(f)/9)):
                lower_bound = c*9
                features.append(np.mean(f[lower_bound:(lower_bound+9)], axis=0))

            features = np.array(features)

        #  TODO: save transforms for predicting.
        # reduce dimensionality
        from sklearn.decomposition import PCA
        pca = PCA(n_components=10)
        out = pca.fit_transform(np.array(features).reshape(len(features), -1))

        return out

    @staticmethod
    def add_steps(lon, lat, step=0.009):
        """
        returns the augmented set of coordinates for all the 9 images adjacent to the center.
        0.009 ~ 1km at the equator.
        """
        new_i, new_j = [], []
        for i, j in zip(lon, lat):
            for a in [-step, 0, step]:
                for b in [-step, 0, step]:
                    new_i.append(i + a)
                    new_j.append(j + b)

        return new_i, new_j


#!/usr/bin/env python

# Jan 12 2014 - latest additions / mutilations by Scott Driscoll
# CuriousInventor.com | https://plus.google.com/+ScottDriscoll
#
# Licensed under the BSD license.  See full license in LICENSE file.
# http://www.lightshowpi.com/
#
# Author: Todd Giles (todd@lightshowpi.com)
# Author: Chris Usey (chris.usey@gmail.com)
# Author: Ryan Jennings
"""Play any audio file and synchronize lights to the music

When executed, this script will play an audio file, as well as turn on and off 8 channels
of lights to the music (via the first 8 GPIO channels on the Rasberry Pi), based upon
music it is playing. Many types of audio files are supported (see decoder.py below), but
it has only been tested with wav and mp3 at the time of this writing.

The timing of the lights turning on and off are controlled based upon the frequency response
of the music being played.  A short segment of the music is analyzed via FFT to get the
frequency response across 8 channels in the audio range.  Each light channel is then turned
on or off based upon whether the amplitude of the frequency response in the corresponding
channel has crossed a dynamic threshold.

The threshold for each channel is "dynamic" in that it is adjusted upwards and downwards
during the song play back based upon the frequency response amplitude of the song. This ensures
that soft songs, or even soft portions of songs will still turn all 8 channels on and off
during the song.

FFT calculation is quite CPU intensive and can adversely affect playback of songs (especially if
attempting to decode the song as well, as is the case for an mp3).  For this reason, the timing
values of the lights turning on and off is cached after it is calculated upon the first time a
new song is played.  The values are cached in a gzip'd text file in the same location as
the song itself.  Subsequent requests to play the same song will use the cached information and not
recompute the FFT, thus reducing CPU utilization dramatically and allowing for clear music
playback of all audio file types.

Sample usage:

sudo python synchronized_lights.py --playlist=/home/pi/music/.playlist
sudo python synchronized_lights.py --file=/home/pi/music/jingle_bells.mp3

Third party dependencies:

alsaaudio: for audio output - http://pyalsaaudio.sourceforge.net/
decoder.py: decoding mp3, ogg, wma, ... - https://pypi.python.org/pypi/decoder.py/1.5XB
numpy: for FFT calcuation - http://www.numpy.org/
"""

import argparse
import csv
import fcntl
import gzip
import logging
import os
import random
from struct import unpack
import sys
import time
import wave
import serial

import alsaaudio as aa
import configuration_manager as cm
import decoder
import hardware_controller as hc
import numpy as np

_PLAYLIST_PATH = "/somepath"

ser = serial.Serial('/dev/ttyACM0',9600)

def setColor(r, color):
	r3 = r%10
	r //= 10
	r2 = r%10
	r //= 10
	r1 = r%10

	ser.write(str(r1).encode())
	ser.write(str(r2).encode())
	ser.write(str(r3).encode())
	ser.write(color.encode())

def colorAllOff():
    setColor(0,'r')
    setColor(0,'g')
    setColor(0,'b')
# RGB LED STRIP STUFF
colorAllOff()

led_array = [0 for i in range(151)]

c = 0.0
columns = [1.0,1.0,1.0,1.0,1.0]
decay = .9
# this writes out light and color information to a continuous RGB LED
# strip that's been wrapped around into 5 columns.
# numbers comes in at 9-15 ish
def display_column(col=0,height=0.0,color=None):
	global c
	global columns
	color = wheel_color(int(c))
	c = c + .1
	if c > 384:
		c = 0.0
	height = height - 9.0
	height = height / 5
	if height < .05:
		height = .05
	elif height > 1.0:
		height = 1.0

	if height < columns[col]:
		columns[col] = columns[col] * decay
		height = columns[col]
	else:
		columns[col] = height

    ##Change here to add my pi support
	if col == 1:
		setColor(int(height*255),'r')

	elif col == 2:
		setColor(int(height*255),'g')

	elif col == 3:
		setColor(int(height*255),'b')

## Configurations - TODO(todd): Move more of this into configuration manager
#_CONFIG = cm.CONFIG
#_LIMIT_LIST = [int(lim) for lim in _CONFIG.get('auto_tuning', 'limit_list').split(',')]
#_LIMIT_THRESHOLD = _CONFIG.getfloat('auto_tuning', 'limit_threshold')
#_LIMIT_THRESHOLD_INCREASE = _CONFIG.getfloat('auto_tuning', 'limit_threshold_increase')
#_LIMIT_THRESHOLD_DECREASE = _CONFIG.getfloat('auto_tuning', 'limit_threshold_decrease')
#_MAX_OFF_CYCLES = _CONFIG.getfloat('auto_tuning', 'max_off_cycles')
_MIN_FREQUENCY = 30
_MAX_FREQUENCY = 16000
#_RANDOMIZE_PLAYLIST = _CONFIG.getboolean('lightshow', 'randomize_playlist')
#try:
    #_CUSTOM_CHANNEL_MAPPING = [int(channel) for channel in
                               #_CONFIG.get('audio_processing', 'custom_channel_mapping').split(',')]
#except:
_CUSTOM_CHANNEL_MAPPING = 0
#try:
    #_CUSTOM_CHANNEL_FREQUENCIES = [int(channel) for channel in
                                   #_CONFIG.get('audio_processing',
                                               #'custom_channel_frequencies').split(',')]
#except:
_CUSTOM_CHANNEL_FREQUENCIES = 0
#try:
    #_PLAYLIST_PATH = _CONFIG.get('lightshow', 'playlist_path').replace('$SYNCHRONIZED_LIGHTS_HOME',
                                                                       #cm.HOME_DIR)
#except:
_PLAYLIST_PATH = "/home/pi/music/.playlist"
CHUNK_SIZE = 2048  # Use a multiple of 8

##What i added

def execute_preshow(config):
    '''Execute the "Preshow" for the given preshow configuration'''
    for transition in config['transitions']:
        start = time.time()
        if transition['type'].lower() == 'on':
            hc.turn_on_lights(True)
        else:
            hc.turn_off_lights(True)
        logging.debug('Transition to ' + transition['type'] + ' for '
            + str(transition['duration']) + ' seconds')
        while transition['duration'] > (time.time() - start):
            cm.load_state()  # Force a refresh of state from file
            play_now = int(cm.get_state('play_now', 0))
            if play_now:
                return  # Skip out on the rest of the preshow

            # Check once every ~ .1 seconds to break out
            time.sleep(0.1)

def calculate_channel_frequency(min_frequency, max_frequency, custom_channel_mapping,
                                custom_channel_frequencies):
	a = 1
    #'''Calculate frequency values for each channel, taking into account custom settings.'''

    ## How many channels do we need to calculate the frequency for
    #if custom_channel_mapping != 0 and len(custom_channel_mapping) == hc.GPIOLEN:
        #logging.debug("Custom Channel Mapping is being used: %s", str(custom_channel_mapping))
        ##channel_length = max(custom_channel_mapping)
    #else:
        #logging.debug("Normal Channel Mapping is being used.")
        ##channel_length = hc.GPIOLEN

    ##logging.debug("Calculating frequencies for %d channels.", channel_length)
    #octaves = (np.log(max_frequency / min_frequency)) / np.log(2)
    #logging.debug("octaves in selected frequency range ... %s", octaves)
    ##octaves_per_channel = octaves / channel_length
    #frequency_limits = []
    #frequency_store = []

    #frequency_limits.append(min_frequency)
    #if custom_channel_frequencies != 0 and (len(custom_channel_frequencies) >= channel_length + 1):
        #logging.debug("Custom channel frequencies are being used")
        #frequency_limits = custom_channel_frequencies
    #else:
        #logging.debug("Custom channel frequencies are not being used")
        #for i in range(1, hc.GPIOLEN + 1):
			#frequency_limits.append(frequency_limits[-1]*2**octaves_per_channel)
            ##frequency_limits.append(frequency_limits[-1]
            ##                        * 10 ** (3 / (10 * (1 / octaves_per_channel))))
    #for i in range(0, channel_length):
        #frequency_store.append((frequency_limits[i], frequency_limits[i + 1]))
        #logging.debug("channel %d is %6.2f to %6.2f ", i, frequency_limits[i],
                      #frequency_limits[i + 1])

    ## we have the frequencies now lets map them if custom mapping is defined
    #if custom_channel_mapping != 0 and len(custom_channel_mapping) == hc.GPIOLEN:
        #frequency_map = []
        #for i in range(0, hc.GPIOLEN):
            #mapped_channel = custom_channel_mapping[i] - 1
            #mapped_frequency_set = frequency_store[mapped_channel]
            #mapped_frequency_set_low = mapped_frequency_set[0]
            #mapped_frequency_set_high = mapped_frequency_set[1]
            #logging.debug("mapped channel: " + str(mapped_channel) + " will hold LOW: "
                          #+ str(mapped_frequency_set_low) + " HIGH: "
                          #+ str(mapped_frequency_set_high))
            #frequency_map.append(mapped_frequency_set)
        #return frequency_map
    #else:
        #return frequency_store

def piff(val, sample_rate):
    '''Return the power array index corresponding to a particular frequency.'''
    return int(CHUNK_SIZE * val / sample_rate)

# TODO(todd): Move FFT related code into separate file as a library
def calculate_levels(data, sample_rate, frequency_limits):
    '''Calculate frequency response for each channel

    Initial FFT code inspired from the code posted here:
    http://www.raspberrypi.org/phpBB3/viewtopic.php?t=35838&p=454041

    Optimizations from work by Scott Driscoll:
    http://www.instructables.com/id/Raspberry-Pi-Spectrum-Analyzer-with-RGB-LED-Strip-/
    '''

    ## create a numpy array. This won't work with a mono file, stereo only.
    #data_stereo = np.frombuffer(data, dtype=np.int16)
    #data = np.empty(len(data) / 4)  # data has two channels and 2 bytes per channel
    #data[:] = data_stereo[::2]  # pull out the even values, just using left channel

    ## if you take an FFT of a chunk of audio, the edges will look like
    ## super high frequency cutoffs. Applying a window tapers the edges
    ## of each end of the chunk down to zero.
    #window = np.hanning(len(data))
    #data = data * window

    ## Apply FFT - real data
    #fourier = np.fft.rfft(data)

    ## Remove last element in array to make it the same size as CHUNK_SIZE
    #fourier = np.delete(fourier, len(fourier) - 1)

    ## Calculate the power spectrum
    #power = np.abs(fourier) ** 2

    #matrix = [0 for i in range(hc.GPIOLEN)]
    #for i in range(hc.GPIOLEN):
        ## take the log10 of the resulting sum to approximate how human ears perceive sound levels
        #matrix[i] = np.log10(np.sum(power[piff(frequency_limits[i][0], sample_rate)
                                          #:piff(frequency_limits[i][1], sample_rate):1]))

    #return matrix

# TODO(todd): Refactor this to make it more readable / modular.
def main():
    '''main'''
    cm = hc.cm

    song_to_play = int(cm.get_state('song_to_play', 0))
    play_now = int(cm.get_state('play_now', 0))

    # Arguments
    parser = argparse.ArgumentParser()
    filegroup = parser.add_mutually_exclusive_group()
    filegroup.add_argument('--playlist', default=_PLAYLIST_PATH,
                           help='Playlist to choose song from.')
    filegroup.add_argument('--file', help='path to the song to play (required if no'
                           'playlist is designated)')
    parser.add_argument('--readcache', type=int, default=1,
                        help='read light timing from cache if available. Default: true')
    args = parser.parse_args()

    # Log everything to our log file
    # TODO(todd): Add logging configuration options.
    #logging.basicConfig(filename=cm.LOG_DIR + '/music_and_lights.play.dbg',
                        #format='[%(asctime)s] %(levelname)s {%(pathname)s:%(lineno)d}'
                        #' - %(message)s',
                        #level=logging.DEBUG)

    # Make sure one of --playlist or --file was specified
    if args.file == None and args.playlist == None:
        print "One of --playlist or --file must be specified"
        sys.exit()

    # Initialize Lights
   # hc.initialize()

    # Only execute preshow if no specific song has been requested to be played right now
    #if not play_now:
        #execute_preshow(cm.lightshow()['preshow'])

    # Determine the next file to play
    song_filename = args.file
    if args.playlist != None and args.file == None:
        most_votes = [None, None, []]
        current_song = None
        with open(args.playlist, 'rb') as playlist_fp:
            fcntl.lockf(playlist_fp, fcntl.LOCK_SH)
            playlist = csv.reader(playlist_fp, delimiter='\t')
            songs = []
            for song in playlist:
                if len(song) < 2 or len(song) > 4:
                    logging.error('Invalid playlist.  Each line should be in the form: '
                                 '<song name><tab><path to song>')
                    sys.exit()
                elif len(song) == 2:
                    song.append(set())
                else:
                    song[2] = set(song[2].split(','))
                    if len(song) == 3 and len(song[2]) >= len(most_votes[2]):
                        most_votes = song
                songs.append(song)
            fcntl.lockf(playlist_fp, fcntl.LOCK_UN)

        if most_votes[0] != None:
            logging.info("Most Votes: " + str(most_votes))
            current_song = most_votes

            # Update playlist with latest votes
            with open(args.playlist, 'wb') as playlist_fp:
                fcntl.lockf(playlist_fp, fcntl.LOCK_EX)
                writer = csv.writer(playlist_fp, delimiter='\t')
                for song in songs:
                    if current_song == song and len(song) == 3:
                        song.append("playing!")
                    if len(song[2]) > 0:
                        song[2] = ",".join(song[2])
                    else:
                        del song[2]
                writer.writerows(songs)
                fcntl.lockf(playlist_fp, fcntl.LOCK_UN)

        else:
            # Get random song
            if _RANDOMIZE_PLAYLIST:
                current_song = songs[random.randint(0, len(songs) - 1)]
            # Get a "play now" requested song
            elif play_now > 0 and play_now <= len(songs):
                current_song = songs[play_now - 1]
            # Play next song in the lineup
            else:
                song_to_play = song_to_play if (song_to_play <= len(songs) - 1) else 0
                current_song = songs[song_to_play]
                next_song = (song_to_play + 1) if ((song_to_play + 1) <= len(songs) - 1) else 0
                cm.update_state('song_to_play', next_song)

        # Get filename to play and store the current song playing in state cfg
        song_filename = current_song[1]
        cm.update_state('current_song', songs.index(current_song))

    song_filename = song_filename.replace("$SYNCHRONIZED_LIGHTS_HOME", '/home/pi/lightshowpi')

    # Ensure play_now is reset before beginning playback
    if play_now:
        cm.update_state('play_now', 0)
        play_now = 0

    # Initialize FFT stats
    #matrix = [0 for _ in range(hc.GPIOLEN)]
    #offct = [0 for _ in range(hc.GPIOLEN)]

    # Build the limit list
    #if len(_LIMIT_LIST) == 1:
    #    limit = [_LIMIT_LIST[0] for _ in range(hc.GPIOLEN)]
    #else:
    #    limit = _LIMIT_LIST

    # Set up audio
    if song_filename.endswith('.wav'):
        musicfile = wave.open(song_filename, 'r')
    else:
        musicfile = decoder.open(song_filename)

    sample_rate = musicfile.getframerate()
    num_channels = musicfile.getnchannels()
    output = aa.PCM(aa.PCM_PLAYBACK, aa.PCM_NORMAL)
    output.setchannels(num_channels)
    output.setrate(sample_rate)
    output.setformat(aa.PCM_FORMAT_S16_LE)
    output.setperiodsize(CHUNK_SIZE)

    # Output a bit about what we're about to play
    song_filename = os.path.abspath(song_filename)
    logging.info("Playing: " + song_filename + " (" + str(musicfile.getnframes() / sample_rate)
                 + " sec)")

    cache = []
    cache_found = False
    cache_filename = os.path.dirname(song_filename) + "/." + os.path.basename(song_filename) \
        + ".sync.gz"
    # The values 12 and 1.5 are good estimates for first time playing back (i.e. before we have
    # the actual mean and standard deviations calculated for each channel).
    #mean = [12.0 for _ in range(hc.GPIOLEN)]
    #std = [1.5 for _ in range(hc.GPIOLEN)]
    if args.readcache:
        # Read in cached fft
        try:
            with gzip.open(cache_filename, 'rb') as playlist_fp:
                cachefile = csv.reader(playlist_fp, delimiter=',')
                for row in cachefile:
                    cache.append([0.0 if np.isinf(float(item)) else float(item) for item in row])
                cache_found = True
                # TODO(todd): Optimize this and / or cache it to avoid delay here
                cache_matrix = np.array(cache)
                for i in range(0, hc.GPIOLEN):
                    std[i] = np.std([item for item in cache_matrix[:, i] if item > 0])
                    mean[i] = np.mean([item for item in cache_matrix[:, i] if item > 0])
                logging.debug("std: " + str(std) + ", mean: " + str(mean))
        except IOError:
            logging.warn("Cached sync data song_filename not found: '" + cache_filename
                         + ".  One will be generated.")

    # Process audio song_filename
    row = 0
    data = musicfile.readframes(CHUNK_SIZE)
    frequency_limits = calculate_channel_frequency(_MIN_FREQUENCY,
                                                   _MAX_FREQUENCY,
                                                   _CUSTOM_CHANNEL_MAPPING,
                                                   _CUSTOM_CHANNEL_FREQUENCIES)

    while data != '' and not play_now:
        output.write(data)

        # Control lights with cached timing values if they exist
        matrix = None
        if cache_found and args.readcache:
            if row < len(cache):
                matrix = cache[row]
            else:
                logging.warning("Ran out of cached FFT values, will update the cache.")
                cache_found = False

        if matrix == None:
            # No cache - Compute FFT in this chunk, and cache results
            matrix = calculate_levels(data, sample_rate, frequency_limits)
            cache.append(matrix)

        # blank out the display
        colorAllOff()
        #for i in range(0, hc.GPIOLEN):
            #if hc.is_pin_pwm(i):
                ## Output pwm, where off is at 0.5 std below the mean
                ## and full on is at 0.75 std above the mean.

                #display_column(i,matrix[i])

                ##brightness = matrix[i] - mean[i] + 0.5 * std[i]
                ##brightness = brightness / (1.25 * std[i])
                ##if brightness > 1.0:
                    ##brightness = 1.0
                ##if brightness < 0:
                    ##brightness = 0
                ##hc.turn_on_light(i, True, int(brightness * 60))
            #else:
                #if limit[i] < matrix[i] * _LIMIT_THRESHOLD:
                    #limit[i] = limit[i] * _LIMIT_THRESHOLD_INCREASE
                    #logging.debug("++++ channel: {0}; limit: {1:.3f}".format(i, limit[i]))
                ## Amplitude has reached threshold
                #if matrix[i] > limit[i]:
                    #hc.turn_on_light(i, True)
                    #offct[i] = 0
                #else:  # Amplitude did not reach threshold
                    #offct[i] = offct[i] + 1
                    #if offct[i] > _MAX_OFF_CYCLES:
                        #offct[i] = 0
                        #limit[i] = limit[i] * _LIMIT_THRESHOLD_DECREASE  # old value 0.8
                    #logging.debug("---- channel: {0}; limit: {1:.3f}".format(i, limit[i]))
                    #hc.turn_off_light(i, True)

        # send out data to RGB LED Strip
        #led.update()
        # Read next chunk of data from music song_filename
        data = musicfile.readframes(CHUNK_SIZE)
        row = row + 1

        # Load new application state in case we've been interrupted
        cm.load_state()
        play_now = int(cm.get_state('play_now', 0))

    if not cache_found:
        with gzip.open(cache_filename, 'wb') as playlist_fp:
            writer = csv.writer(playlist_fp, delimiter=',')
            writer.writerows(cache)
            logging.info("Cached sync data written to '." + cache_filename
                         + "' [" + str(len(cache)) + " rows]")

    # We're done, turn it all off ;)
    hc.clean_up()

if __name__ == "__main__":
    main()