import threading
import pygame
import math
import logging
import numpy
import scipy.io as sio

class Game(object):
    """Docstring for Game """
    G_LEFT=1
    G_RIGHT=2
    G_JUMP=3

    def __init__(self,cnf):
       """@todo: to be defined """
       #list of moves for the stage
       self.goal=None
       self.trials=[]
       self.classifications=[]
       self.currentClassification=[]
       self.view=None
       self.clazzes=[]
       self.logger=logging.getLogger("logger")
       self.finished=False
       self.cnf=cnf
       self.onTrial=False
       self.sound= pygame.mixer.Sound(cnf['snd'])

    def setView(self,view):
            self.view=view

    def startTrial(self,bus,clazz):
        if not self.finished:
                self.view.prep=True
                self.onTrial=True

    def startLiveClassification(self,bus,clazz):
        if not self.finished:
                self.view.prep=False

                self.logger.debug('play start!')
                self.sound.play()
                self.logger.debug('play stop!')

                #compute the goal having the zombie position+class
                if clazz==3:#will have to stay in the same lane
                   goal=Game.G_JUMP
                elif clazz==1:
                    goal=Game.G_LEFT
                elif clazz==2:
                    goal=Game.G_RIGHT
                self.trials.append(clazz)
                #update the view
                self.view.setGoal(goal)

    def endTrial(self,bus,clazz):
        if not self.finished:
                self.onTrial=False
                #update trials and classification result
                self.classifications.append(self.currentClassification[-1])
                #reset goal and the current classificaiton
                self.goal=None
                #if not self.checkStack():
                self.view.zombie.move(self.currentClassification[-1])
                self.currentClassification=[]

                #update view
                self.view.setGoal(None)

    def checkStack(self):
            """@todo: Docstring for checkStack.
            :returns: @todo

            """
            print "check stack"
            print self.trials[-1]!=self.classifications[-1]
            #error in the last
            if  self.trials[-1]!=self.classifications[-1]:
                   print "HEY0"
                   #two errors in a row for the last class
                   if self.view.zombie.pos==2 and self.classifications[-1]==2:
                           print "HEY1"
                           self.view.zombie.move(1)
                           return True

                   elif self.view.zombie.pos==0 and self.classifications[-1]==1:
                           #move to center
                           self.view.zombie.move(2)
                           return True
            return False

    def classEvent(self,buz,clazz):
        #update the current classificaiton
        if self.view.goal.goal!=None and clazz!=-1:
                self.currentClassification.append(clazz)
        #self.view.update()

    def theEnd(self,buzz,whatever):
        self.finished=True
        self.logger.info("This is the end")
        self.saveResults()

    def getScore(self):
        hits=0
        for exp,res in zip(self.trials,self.classifications):
                if exp==res:
                        hits+=1
        return hits

    def saveResults(self):
        path=self.cnf['outputpath']+"/online_outs/%s_%s.mat"%(self.cnf['user'],self.cnf['session'])
        ypred=numpy.array(self.classifications)
	sio.savemat(path,{'ypred':ypred})

class GameView(object):
        """Docstring for GameView """

        def __init__(self,game,screen,cnf):
                self.cnf=cnf
                self.game=game
                self.zombie=Zombie(cnf,8)
                self.terrain=Terrain(cnf,20)
                self.goal=Goal(cnf,4,self.zombie)
                self.screen=screen
                self.lock=threading.Lock()
                self.state=GameViewState()
                self.state.bgcolor=cnf['game']['bg']
                self.font=cnf['game']['font']
                self.prep=False
                self.clazz=-1
                self.__load_images()

        def setGoal(self,clazz):
                self.lock.acquire()
                self.clazz=clazz
                if clazz==Game.G_JUMP:
                        self.goal.setGoal(self.zombie.pos)
                elif clazz==Game.G_LEFT:
                        self.goal.setGoal(self.zombie.pos-1)
                elif clazz==Game.G_RIGHT:
                        self.goal.setGoal(self.zombie.pos+1)
                else:
                        self.goal.setGoal(None)
                self.lock.release()

        def onUpdate(self):
                #paint
                #backgound

     #           print "GAME: %s"%threading.currentThread()
                self.lock.acquire()
                if self.game.finished:
                        self.showScore()

                else:
                        self.screen.fill(self.state.bgcolor)
                        if self.prep:
                                self.drawStatusBarPrep()
                        else:
                                self.drawStatusBar()
                        self.terrain.update(self.screen)
                        self.goal.update(self.screen)
                        self.zombie.update(self.screen)
                self.lock.release()

        def showScore(self):
                self.screen.fill(self.state.bgcolor)
                font = pygame.font.Font(self.font, 120)
                #trials msg
                surface = font.render("Game Over!", True, (10, 10, 10))
                rect = surface.get_rect()
                rect.center=(750,200)
                self.screen.blit(surface, rect)

                surface = font.render("Score: %i "%(self.game.getScore()*100), True, (10, 10, 10))
                rect = surface.get_rect()
                rect.center=(750,300)
                self.screen.blit(surface, rect)

                surface = font.render("Hit ratio: %2.1f %% "%(100*float(self.game.getScore())/len(self.game.trials)), True, (10, 10, 10))
                rect = surface.get_rect()
                rect.center=(750,350)
                self.screen.blit(surface, rect)

        def drawStatusBar(self):
                pygame.draw.rect(self.screen, self.cnf['game']['st'], (0,800,1500,200), 0)
                font = pygame.font.Font(self.font, 80)
                #trials msg
                surface = font.render("Trials: %i"%len(self.game.trials), True, (10, 10, 10))
                rect = surface.get_rect()
                rect.center=(1100,900)
                self.screen.blit(surface, rect)
                #score
                surface = font.render("Score: %i"%(self.game.getScore()*100), True, (10, 10, 10))
                rect = surface.get_rect()
                rect.center=(1300,900)
                self.screen.blit(surface, rect)
                #arrows
                imgs=[self.arrow]*3
                if self.goal.goal!=None and self.clazz!=None:
                        imgs[self.clazz-1]=self.arrow_class
                self.drawArrow(1,imgs[0].copy())
                self.drawArrow(2,imgs[1].copy())
                self.drawArrow(3,imgs[2].copy())

        def __load_images(self):
                img = pygame.image.load(self.cnf['game']['arrow']).convert_alpha()
                imgClass=pygame.image.load(self.cnf['game']['arrow_class']).convert_alpha()
                img = pygame.transform.scale(img,(100,100))
                imgClass = pygame.transform.scale(imgClass,(120,120))
                self.arrow=img
                self.arrow_class=imgClass
                self.cross=pygame.image.load("imgs/cross.png").convert_alpha()

        def drawStatusBarPrep(self):
                pygame.draw.rect(self.screen, self.cnf['game']['st'], (0,800,1500,200), 0)
                font = pygame.font.Font(self.font, 80)
                img = self.cross

                rect = img.get_rect()
                rect.center=(750,900-5)
                self.screen.blit(img, rect)

        def drawArrow(self,direct,img):
                pixarr=pygame.PixelArray(img)
                color=(0,0,0)
                if len(self.game.currentClassification)>0:
                        if self.game.currentClassification[-1]==direct:
                                color=(215,46,46)
                elif len(self.game.classifications)>0 and self.game.classifications[-1]==direct and not self.game.onTrial:
                                color=(28,220,28)
                pixarr.replace((0,255,0),color,distance=0)
                img=pixarr.surface
                del pixarr
                rect = img.get_rect()
                rect2 = img.copy().get_rect()
                if direct==Game.G_RIGHT:#right
                        rect.center = (750+150,900)
                elif direct==Game.G_LEFT:#left
                        img=pygame.transform.rotate(img,180)
                        rect.center = (750-150,900)
                else:
                        img=pygame.transform.rotate(img,90)
                        rect.center = (750,900)

                self.screen.blit(img,rect)

class GameViewState:

        def __init__(self):
                self.bgcolor=None

class Goal(object):
        """Docstring for Goal """

        def __init__(self,cnf,updateRate,zombie):
                self.cnf=cnf
                self.pics=[]
                self.tick=0
                self.step=0
                self.goal=None
                self.pos=[0,0]
                self.updateRate=updateRate
                self.zombie=zombie
                self.totalSteps=160
                self.stepSize=[0,(750.)/self.totalSteps]
                self.angle=0.0
                self.bomb=True
                self.lock=threading.Lock()
                self.__load_pics()

        def __load_pics(self):
                imgs=self.cnf['game']['goals']
                for idx,img in enumerate(imgs):
                       self.pics.append([])
                       pic=pygame.image.load(img).convert_alpha()
                       for step in range(self.totalSteps+1):
                               size=int(20.+step/3)
                               scaled = pygame.transform.scale(pic.copy(),(size,size))
                               self.pics[idx].append(scaled)

        def setGoal(self,goal):
                self.lock.acquire()
                if goal == None:
                        self.goal==None
                elif self.goal==None:
                        self.goal=goal
                        self.step=0
                        if goal==self.zombie.pos:
                                self.img=1
                                self.bomb=True
                                iniY=30
                                self.totalSteps=130
                                self.stepSize=[0,(700.)/self.totalSteps]
                        else:
                                self.img=0
                                self.bomb=False
                                iniY=0
                                self.totalSteps=130
                                self.stepSize=[0,(700.)/self.totalSteps]

                        self.pos=[750+(self.goal-1)*55,iniY]
                        self.angle=(self.goal-1)*0.20

                self.lock.release()

        def update(self,screen):
               if self.goal==None:
                       return
               self.lock.acquire()

               if self.tick % self.updateRate == 0:
                      self.step+=1
                      self.pos[0]+=self.stepSize[0]
                      self.pos[1]+=self.stepSize[1]

               if not self.bomb:
                       self.drawGoal(screen,self.pos,self.angle)
               else:
                       pos=[750-55,self.pos[1]]
                       self.drawGoal(screen,pos,-0.20)
                       pos=[750,self.pos[1]]
                       self.drawGoal(screen,pos,0)
                       pos=[750+55,self.pos[1]]
                       self.drawGoal(screen,pos,0.20)

               if self.step==self.totalSteps:
                       self.goal=None

               self.tick+=1
               self.tick=self.tick % self.updateRate
               self.lock.release()

        def drawGoal(self,screen,pos,angle):
               img = self.pics[self.img][self.step]
               rect = img.get_rect()
               rect.center = (pos[0]+(angle*(pos[1]-10)),pos[1])
               screen.blit(img,rect)

class Terrain(object):
        """Docstring for Terrain """

        def __init__(self,cnf,updateRate):
                """@todo: to be defined """
                self.cnf=cnf
                self.pics=[]
                self.tick=0
                self.step=0
                self.updateRate=updateRate
                self.__load_pics()

        def __load_pics(self):
                imgs=self.cnf['game']['terrain']
                for img in imgs:
                       pic=pygame.image.load(img).convert_alpha()
                       self.pics.append(pic)

        def update(self,screen):
               if self.tick % self.updateRate == 0:
                      self.step+=1

               self.step=self.step % len(self.pics)

               img=self.pics[self.step]#.copy()
               rect = img.get_rect()
               rect.x= 0
               rect.y= 0
               screen.blit(img,rect)
               self.tick+=1
               self.tick=self.tick % self.updateRate

class Zombie(object):
        """Docstring for Zombie """

        def __init__(self,cnf,updateRate):
                """@todo: to be defined """
                self.cnf=cnf
                self.pos=1
                self.nextPos=1
                self.curPos=750
                self.pics=[]
                self.jumpingPics=[]
                self.jump=False
                self.jumpStep=0
                self.tick=0
                self.updateRate=updateRate
                self.step=0
                self.__load_pics()

        def __load_pics(self):
                imgs=self.cnf['game']['c_walk']
                for img in imgs:
                       pic=pygame.image.load(img).convert_alpha()
                       pic=pygame.transform.scale(pic,(30,90))
                       self.pics.append(pic)

                imgs=self.cnf['game']['c_jump']
                for img in imgs:
                       pic=pygame.image.load(img).convert_alpha()
                       pic=pygame.transform.scale(pic,(30,114))
                       self.jumpingPics.append(pic)

        def __str__(self):
                return "[zombie curPos:%i pos:%i next:%i]"%(self.curPos,self.pos,self.nextPos)

        def move(self,mov):
                if mov==Game.G_LEFT and self.isStill():
                        if self.pos>0:
                                self.nextPos-=1
                elif mov==Game.G_RIGHT and self.isStill():
                        if self.pos<2:
                                self.nextPos+=1
                elif mov==Game.G_JUMP:
                        self.jump=True
                        self.jumpStep=1.

        def isStill(self):
                return self.curPos==750 or self.curPos==750+200 or self.curPos ==750-200

        def __move(self):
                if self.nextPos!=self.pos:
                       self.curPos+=(self.nextPos-self.pos)*5
                       if self.isStill():
                                self.pos=self.nextPos
                                if self.pos==0 or self.pos==2:
                                        self.nextPos=1

                elif self.jump:
                        pass
                        self.jumpStep+=0.065
                        if self.jumpStep>=len(self.jumpingPics):
                                self.jump=False

        def update(self,screen):
               if self.tick % self.updateRate == 0:
                      self.step+=1
               self.step=self.step % len(self.pics)
               img=None
               yAdj=0
               if not self.jump:
                        img=self.pics[self.step].copy()

               else:
                       yAdj=4
                       idx=int(math.floor(self.jumpStep))
                       #if idx>len(self.jumpingPics):
                               #idx=len(self.jumpingPics)-idx%len(self.jumpingPics)
                       img=self.jumpingPics[idx-1].copy()

               rect = img.get_rect()
               self.__move()
               rect.center = (self.curPos,700-yAdj)
               screen.blit(img,rect)
               self.tick+=1
               self.tick=self.tick % self.updateRate


# -*- coding: utf-8 -*-
# --------------------------------------------------------
# Implementation-CVPR2015-CNN-for-ReID
# Copyright (c) 2017 Ning Ding
# Licensed under The MIT License [see LICENSE for details]
# Written by Ning Ding
# --------------------------------------------------------

"""
Data Preparation Script.
"""
import numpy as np
np.random.seed(1217)
import h5py
from PIL import Image
from keras.preprocessing import image as pre_image

#Reduce the batch size(ex: 50) if there is a resource exhaust error.
BATCH_SIZE = 150

class NumpyArrayIterator_for_CUHK03(pre_image.Iterator):

    def __init__(self, f, train_or_validation = 'train', flag = 1, image_data_generator = None,
                 batch_size = BATCH_SIZE, shuffle=True, seed=1217):
        self.f = f
        self.length = len(f['a'][train_or_validation].keys())
        self.train_or_validation = train_or_validation
        self.flag = flag
        self.image_data_generator = image_data_generator
        super(NumpyArrayIterator_for_CUHK03, self).__init__(3000000, batch_size / 2, shuffle, seed)

    def next(self):
        with self.lock:
            index_array, current_index, current_batch_size = next(self.index_generator)

        batch_x1 = np.zeros(tuple([current_batch_size * 2] + [160,60,3]))
        batch_x2 = np.zeros(tuple([current_batch_size * 2] + [160,60,3]))
        batch_y  = np.zeros([current_batch_size * 2, 2])

        for i, j in enumerate(index_array):

            k = np.random.randint(self.length)
            while k == 1155:
                k = np.random.randint(self.length)
            ja = np.random.randint(self.f['a'][self.train_or_validation][str(k)].shape[0])
            jb = np.random.randint(self.f['b'][self.train_or_validation][str(k)].shape[0])

            x1 = self.f['a'][self.train_or_validation][str(k)][ja]
            x2 = self.f['b'][self.train_or_validation][str(k)][jb]
            if np.random.rand() > self.flag:
                x1 = self.image_data_generator.random_transform(x1.astype('float32'))
            if np.random.rand() > self.flag:
                x2 = self.image_data_generator.random_transform(x2.astype('float32'))

            batch_x1[2*i] = x1
            batch_x2[2*i] = x2
            batch_y[2*i][1] = 1

            ka,kb = np.random.choice(range(self.length),2)
            while ka == 1155 or kb == 1155:
                ka,kb = np.random.choice(range(self.length),2)

            ja = np.random.randint(self.f['a'][self.train_or_validation][str(ka)].shape[0])
            jb = np.random.randint(self.f['b'][self.train_or_validation][str(kb)].shape[0])

            x1 = self.f['a'][self.train_or_validation][str(ka)][ja]
            x2 = self.f['b'][self.train_or_validation][str(kb)][jb]

            batch_x1[2*i+1] = x1
            batch_x2[2*i+1] = x2
            batch_y[2*i+1][0] = 1

        return [batch_x1,batch_x2], batch_y

class ImageDataGenerator_for_multiinput(pre_image.ImageDataGenerator):

    def flow(self, f, train_or_validation = 'train', flag = 0, batch_size = BATCH_SIZE, shuffle=True, seed=1217):

        return NumpyArrayIterator_for_CUHK03(f, train_or_validation, flag, self, batch_size=batch_size, shuffle=shuffle, seed=seed)

    def agumentation(self, X, rounds=1, seed=None):

        if seed is not None:
            np.random.seed(seed)

        X = np.copy(X)
        aX = np.zeros(tuple([rounds * X.shape[0]] + list(X.shape)[1:]))
        for r in range(rounds):
            for i in range(X.shape[0]):
                aX[i + r * X.shape[0]] = self.random_transform(X[i])
        X = aX
        return X


from datasets.nuScenes.nuScenes import NuScenesTrajectories
from nuscenes.prediction.input_representation.static_layers import correct_yaw
from nuscenes.map_expansion.map_api import NuScenesMap
from nuscenes.eval.common.utils import quaternion_yaw
from pyquaternion import Quaternion
from nuscenes.prediction import PredictHelper
import numpy as np
from typing import Dict, Tuple, Union, List
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon
import os
import pickle
import torch

class NuScenesVector(NuScenesTrajectories):
    """
    NuScenes dataset class for single agent prediction, using the vector representation for maps and agents
    """

    def __init__(self, mode: str, data_dir: str, args: Dict, helper: PredictHelper):
        """
        Initialize predict helper, agent and scene representations
        :param mode: Mode of operation of dataset, one of {'compute_stats', 'extract_data', 'load_data'}
        :param data_dir: Directory to store extracted pre-processed data
        :param helper: NuScenes PredictHelper
        :param args: Dataset arguments
        """
        super().__init__(mode, data_dir, args, helper)

        # Initialize helper and maps
        self.map_locs = ['singapore-onenorth', 'singapore-hollandvillage', 'singapore-queenstown', 'boston-seaport']
        self.maps = {i: NuScenesMap(map_name=i, dataroot=self.helper.data.dataroot) for i in self.map_locs}

        # Vector map parameters
        self.map_extent = args['map_extent']
        self.polyline_resolution = args['polyline_resolution']
        self.polyline_length = args['polyline_length']

        # Load dataset stats (max nodes, max agents etc.)
        if self.mode == 'extract_data':
            stats = self.load_stats()
            self.max_nodes = stats['num_lane_nodes']
            self.max_vehicles = stats['num_vehicles']
            self.max_pedestrians = stats['num_pedestrians']

        # Whether to add random flips for data augmentation
        elif self.mode == 'load_data':
            self.random_flips = args['random_flips']

    def compute_stats(self, idx: int) -> Dict[str, int]:
        """
        Function to compute statistics for a given data point
        """
        num_lane_nodes = self.get_map_representation(idx)
        num_vehicles, num_pedestrians = self.get_surrounding_agent_representation(idx)
        stats = {
            'num_lane_nodes': num_lane_nodes,
            'num_vehicles': num_vehicles,
            'num_pedestrians': num_pedestrians
        }

        return stats

    def load_data(self, idx: int) -> Dict:
        """
        Perform random flips if lag is set to true.
        """
        data = super().load_data(idx)

        if self.random_flips:
            if torch.randint(2, (1, 1)).squeeze().bool().item():
                data = self.flip_horizontal(data)

        return data

    def get_target_agent_representation(self, idx: int) -> np.ndarray:
        """
        Extracts target agent representation
        :param idx: data index
        :return hist: track history for target agent, shape: [t_h * 2, 5]
        """
        i_t, s_t = self.token_list[idx].split("_")

        # x, y co-ordinates in agent's frame of reference
        hist = self.helper.get_past_for_agent(i_t, s_t, seconds=self.t_h, in_agent_frame=True)

        # Zero pad for track histories shorter than t_h
        hist_zeropadded = np.zeros((int(self.t_h) * 2 + 1, 2))

        # Flip to have correct order of timestamps
        hist = np.flip(hist, 0)
        hist_zeropadded[-hist.shape[0]-1: -1] = hist
        hist = hist_zeropadded

        # Get velocity, acc and yaw_rate over past t_h sec
        motion_states = self.get_past_motion_states(i_t, s_t)
        hist = np.concatenate((hist, motion_states), axis=1)

        return hist

    def get_map_representation(self, idx: int) -> Union[int, Dict]:
        """
        Extracts map representation
        :param idx: data index
        :return: Returns an ndarray with lane node features, shape [max_nodes, polyline_length, 5] and an ndarray of
            masks of the same shape, with value 1 if the nodes/poses are empty,
        """
        i_t, s_t = self.token_list[idx].split("_")
        map_name = self.helper.get_map_name_from_sample_token(s_t)
        map_api = self.maps[map_name]

        # Get agent representation in global co-ordinates
        global_pose = self.get_target_agent_global_pose(idx)

        # Get lanes around agent within map_extent
        lanes = self.get_lanes_around_agent(global_pose, map_api)

        # Get relevant polygon layers from the map_api
        polygons = self.get_polygons_around_agent(global_pose, map_api)

        # Get vectorized representation of lanes
        lane_node_feats, _ = self.get_lane_node_feats(global_pose, lanes, polygons)

        # Discard lanes outside map extent
        lane_node_feats = self.discard_poses_outside_extent(lane_node_feats)

        # Add dummy node (0, 0, 0, 0, 0) if no lane nodes are found
        if len(lane_node_feats) == 0:
            lane_node_feats = [np.zeros((1, 5))]

        # While running the dataset class in 'compute_stats' mode:
        if self.mode == 'compute_stats':
            return len(lane_node_feats)

        # Convert list of lane node feats to fixed size numpy array and masks
        lane_node_feats, lane_node_masks = self.list_to_tensor(lane_node_feats, self.max_nodes, self.polyline_length, 5)

        map_representation = {
            'lane_node_feats': lane_node_feats,
            'lane_node_masks': lane_node_masks
        }

        return map_representation

    def get_surrounding_agent_representation(self, idx: int) -> \
            Union[Tuple[int, int], Dict]:
        """
        Extracts surrounding agent representation
        :param idx: data index
        :return: ndarrays with surrounding pedestrian and vehicle track histories and masks for non-existent agents
        """

        # Get vehicles and pedestrian histories for current sample
        vehicles = self.get_agents_of_type(idx, 'vehicle')
        pedestrians = self.get_agents_of_type(idx, 'human')

        # Discard poses outside map extent
        vehicles = self.discard_poses_outside_extent(vehicles)
        pedestrians = self.discard_poses_outside_extent(pedestrians)

        # While running the dataset class in 'compute_stats' mode:
        if self.mode == 'compute_stats':
            return len(vehicles), len(pedestrians)

        # Convert to fixed size arrays for batching
        vehicles, vehicle_masks = self.list_to_tensor(vehicles, self.max_vehicles, self.t_h * 2 + 1, 5)
        pedestrians, pedestrian_masks = self.list_to_tensor(pedestrians, self.max_pedestrians, self.t_h * 2 + 1, 5)

        surrounding_agent_representation = {
            'vehicles': vehicles,
            'vehicle_masks': vehicle_masks,
            'pedestrians': pedestrians,
            'pedestrian_masks': pedestrian_masks
        }

        return surrounding_agent_representation

    def get_target_agent_global_pose(self, idx: int) -> Tuple[float, float, float]:
        """
        Returns global pose of target agent
        :param idx: data index
        :return global_pose: (x, y, yaw) or target agent in global co-ordinates
        """
        i_t, s_t = self.token_list[idx].split("_")
        sample_annotation = self.helper.get_sample_annotation(i_t, s_t)
        x, y = sample_annotation['translation'][:2]
        yaw = quaternion_yaw(Quaternion(sample_annotation['rotation']))
        yaw = correct_yaw(yaw)
        global_pose = (x, y, yaw)

        return global_pose

    def get_lanes_around_agent(self, global_pose: Tuple[float, float, float], map_api: NuScenesMap) -> Dict:
        """
        Gets lane polylines around the target agent
        :param global_pose: (x, y, yaw) or target agent in global co-ordinates
        :param map_api: nuScenes map api
        :return lanes: Dictionary of lane polylines
        """
        x, y, _ = global_pose
        radius = max(self.map_extent)
        lanes = map_api.get_records_in_radius(x, y, radius, ['lane', 'lane_connector'])
        lanes = lanes['lane'] + lanes['lane_connector']
        lanes = map_api.discretize_lanes(lanes, self.polyline_resolution)

        return lanes

    def get_polygons_around_agent(self, global_pose: Tuple[float, float, float], map_api: NuScenesMap) -> Dict:
        """
        Gets polygon layers around the target agent e.g. crosswalks, stop lines
        :param global_pose: (x, y, yaw) or target agent in global co-ordinates
        :param map_api: nuScenes map api
        :return polygons: Dictionary of polygon layers, each type as a list of shapely Polygons
        """
        x, y, _ = global_pose
        radius = max(self.map_extent)
        record_tokens = map_api.get_records_in_radius(x, y, radius, ['stop_line', 'ped_crossing'])
        polygons = {k: [] for k in record_tokens.keys()}
        for k, v in record_tokens.items():
            for record_token in v:
                polygon_token = map_api.get(k, record_token)['polygon_token']
                polygons[k].append(map_api.extract_polygon(polygon_token))

        return polygons

    def get_lane_node_feats(self, origin: Tuple, lanes: Dict[str, List[Tuple]],
                            polygons: Dict[str, List[Polygon]]) -> Tuple[List[np.ndarray], List[str]]:
        """
        Generates vector HD map representation in the agent centric frame of reference
        :param origin: (x, y, yaw) of target agent in global co-ordinates
        :param lanes: lane centerline poses in global co-ordinates
        :param polygons: stop-line and cross-walk polygons in global co-ordinates
        :return:
        """

        # Convert lanes to list
        lane_ids = [k for k, v in lanes.items()]
        lanes = [v for k, v in lanes.items()]

        # Get flags indicating whether a lane lies on stop lines or crosswalks
        lane_flags = self.get_lane_flags(lanes, polygons)

        # Convert lane polylines to local coordinates:
        lanes = [np.asarray([self.global_to_local(origin, pose) for pose in lane]) for lane in lanes]

        # Concatenate lane poses and lane flags
        lane_node_feats = [np.concatenate((lanes[i], lane_flags[i]), axis=1) for i in range(len(lanes))]

        # Split lane centerlines into smaller segments:
        lane_node_feats, lane_node_ids = self.split_lanes(lane_node_feats, self.polyline_length, lane_ids)

        return lane_node_feats, lane_node_ids

    def get_agents_of_type(self, idx: int, agent_type: str) -> List[np.ndarray]:
        """
        Returns surrounding agents of a particular class for a given sample
        :param idx: data index
        :param agent_type: 'human' or 'vehicle'
        :return: list of ndarrays of agent track histories.
        """
        i_t, s_t = self.token_list[idx].split("_")

        # Get agent representation in global co-ordinates
        origin = self.get_target_agent_global_pose(idx)

        # Load all agents for sample
        agent_details = self.helper.get_past_for_sample(s_t, seconds=self.t_h, in_agent_frame=False, just_xy=False)
        agent_hist = self.helper.get_past_for_sample(s_t, seconds=self.t_h, in_agent_frame=False, just_xy=True)

        # Add present time to agent histories
        present_time = self.helper.get_annotations_for_sample(s_t)
        for annotation in present_time:
            ann_i_t = annotation['instance_token']
            if ann_i_t in agent_hist.keys():
                present_pose = np.asarray(annotation['translation'][0:2]).reshape(1, 2)
                if agent_hist[ann_i_t].any():
                    agent_hist[ann_i_t] = np.concatenate((present_pose, agent_hist[ann_i_t]))
                else:
                    agent_hist[ann_i_t] = present_pose

        # Filter for agent type
        agent_list = []
        agent_i_ts = []
        for k, v in agent_details.items():
            if v and agent_type in v[0]['category_name'] and v[0]['instance_token'] != i_t:
                agent_list.append(agent_hist[k])
                agent_i_ts.append(v[0]['instance_token'])

        # Convert to target agent's frame of reference
        for agent in agent_list:
            for n, pose in enumerate(agent):
                local_pose = self.global_to_local(origin, (pose[0], pose[1], 0))
                agent[n] = np.asarray([local_pose[0], local_pose[1]])

        # Flip history to have most recent time stamp last and extract past motion states
        for n, agent in enumerate(agent_list):
            xy = np.flip(agent, axis=0)
            motion_states = self.get_past_motion_states(agent_i_ts[n], s_t)
            motion_states = motion_states[-len(xy):, :]
            agent_list[n] = np.concatenate((xy, motion_states), axis=1)

        return agent_list

    def discard_poses_outside_extent(self, pose_set: List[np.ndarray],
                                     ids: List[str] = None) -> Union[List[np.ndarray],
                                                                     Tuple[List[np.ndarray], List[str]]]:
        """
        Discards lane or agent poses outside predefined extent in target agent's frame of reference.
        :param pose_set: agent or lane polyline poses
        :param ids: annotation record tokens for pose_set. Only applies to lanes.
        :return: Updated pose set
        """
        updated_pose_set = []
        updated_ids = []

        for m, poses in enumerate(pose_set):
            flag = False
            for n, pose in enumerate(poses):
                if self.map_extent[0] <= pose[0] <= self.map_extent[1] and \
                        self.map_extent[2] <= pose[1] <= self.map_extent[3]:
                    flag = True

            if flag:
                updated_pose_set.append(poses)
                if ids is not None:
                    updated_ids.append(ids[m])

        if ids is not None:
            return updated_pose_set, updated_ids
        else:
            return updated_pose_set

    def load_stats(self) -> Dict[str, int]:
        """
        Function to load dataset statistics like max surrounding agents, max nodes, max edges etc.
        """
        filename = os.path.join(self.data_dir, 'stats.pickle')
        if not os.path.isfile(filename):
            raise Exception('Could not find dataset statistics. Please run the dataset in compute_stats mode')

        with open(filename, 'rb') as handle:
            stats = pickle.load(handle)

        return stats

    def get_past_motion_states(self, i_t, s_t):
        """
        Returns past motion states: v, a, yaw_rate for a given instance and sample token over self.t_h seconds
        """
        motion_states = np.zeros((2 * self.t_h + 1, 3))
        motion_states[-1, 0] = self.helper.get_velocity_for_agent(i_t, s_t)
        motion_states[-1, 1] = self.helper.get_acceleration_for_agent(i_t, s_t)
        motion_states[-1, 2] = self.helper.get_heading_change_rate_for_agent(i_t, s_t)
        hist = self.helper.get_past_for_agent(i_t, s_t, seconds=self.t_h, in_agent_frame=True, just_xy=False)

        for k in range(len(hist)):
            motion_states[-(k + 2), 0] = self.helper.get_velocity_for_agent(i_t, hist[k]['sample_token'])
            motion_states[-(k + 2), 1] = self.helper.get_acceleration_for_agent(i_t, hist[k]['sample_token'])
            motion_states[-(k + 2), 2] = self.helper.get_heading_change_rate_for_agent(i_t, hist[k]['sample_token'])

        motion_states = np.nan_to_num(motion_states)
        return motion_states

    @staticmethod
    def global_to_local(origin: Tuple, global_pose: Tuple) -> Tuple:
        """
        Converts pose in global co-ordinates to local co-ordinates.
        :param origin: (x, y, yaw) of origin in global co-ordinates
        :param global_pose: (x, y, yaw) in global co-ordinates
        :return local_pose: (x, y, yaw) in local co-ordinates
        """
        # Unpack
        global_x, global_y, global_yaw = global_pose
        origin_x, origin_y, origin_yaw = origin

        # Translate
        local_x = global_x - origin_x
        local_y = global_y - origin_y

        # Rotate
        global_yaw = correct_yaw(global_yaw)
        theta = np.arctan2(-np.sin(global_yaw-origin_yaw), np.cos(global_yaw-origin_yaw))

        r = np.asarray([[np.cos(np.pi/2 - origin_yaw), np.sin(np.pi/2 - origin_yaw)],
                        [-np.sin(np.pi/2 - origin_yaw), np.cos(np.pi/2 - origin_yaw)]])
        local_x, local_y = np.matmul(r, np.asarray([local_x, local_y]).transpose())

        local_pose = (local_x, local_y, theta)

        return local_pose

    @staticmethod
    def split_lanes(lanes: List[np.ndarray], max_len: int, lane_ids: List[str]) -> Tuple[List[np.ndarray], List[str]]:
        """
        Splits lanes into roughly equal sized smaller segments with defined maximum length
        :param lanes: list of lane poses
        :param max_len: maximum admissible length of polyline
        :param lane_ids: list of lane ID tokens
        :return lane_segments: list of smaller lane segments
                lane_segment_ids: list of lane ID tokens corresponding to original lane that the segment is part of
        """
        lane_segments = []
        lane_segment_ids = []
        for idx, lane in enumerate(lanes):
            n_segments = int(np.ceil(len(lane) / max_len))
            n_poses = int(np.ceil(len(lane) / n_segments))
            for n in range(n_segments):
                lane_segment = lane[n * n_poses: (n+1) * n_poses]
                lane_segments.append(lane_segment)
                lane_segment_ids.append(lane_ids[idx])

        return lane_segments, lane_segment_ids

    @staticmethod
    def get_lane_flags(lanes: List[List[Tuple]], polygons: Dict[str, List[Polygon]]) -> List[np.ndarray]:
        """
        Returns flags indicating whether each pose on lane polylines lies on polygon map layers
        like stop-lines or cross-walks
        :param lanes: list of lane poses
        :param polygons: dictionary of polygon layers
        :return lane_flags: list of ndarrays with flags
        """

        lane_flags = [np.zeros((len(lane), len(polygons.keys()))) for lane in lanes]
        for lane_num, lane in enumerate(lanes):
            for pose_num, pose in enumerate(lane):
                point = Point(pose[0], pose[1])
                for n, k in enumerate(polygons.keys()):
                    polygon_list = polygons[k]
                    for polygon in polygon_list:
                        if polygon.contains(point):
                            lane_flags[lane_num][pose_num][n] = 1
                            break

        return lane_flags

    @staticmethod
    def list_to_tensor(feat_list: List[np.ndarray], max_num: int, max_len: int,
                       feat_size: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Converts a list of sequential features (e.g. lane polylines or agent history) to fixed size numpy arrays for
        forming mini-batches

        :param feat_list: List of sequential features
        :param max_num: Maximum number of sequences in List
        :param max_len: Maximum length of each sequence
        :param feat_size: Feature dimension
        :return: 1) ndarray of features of shape [max_num, max_len, feat_dim]. Has zeros where elements are missing,
            2) ndarray of binary masks of shape [max_num, max_len, feat_dim]. Has ones where elements are missing.
        """
        feat_array = np.zeros((max_num, max_len, feat_size))
        mask_array = np.ones((max_num, max_len, feat_size))
        for n, feats in enumerate(feat_list):
            feat_array[n, :len(feats), :] = feats
            mask_array[n, :len(feats), :] = 0

        return feat_array, mask_array

    @staticmethod
    def flip_horizontal(data: Dict):
        """
        Helper function to randomly flip some samples across y-axis for data augmentation
        :param data: Dictionary with inputs and ground truth values.
        :return: data: Dictionary with inputs and ground truth values fligpped along y-axis.
        """
        # Flip target agent
        hist = data['inputs']['target_agent_representation']
        hist[:, 0] = -hist[:, 0]  # x-coord
        hist[:, 4] = -hist[:, 4]  # yaw-rate
        data['inputs']['target_agent_representation'] = hist

        # Flip lane node features
        lf = data['inputs']['map_representation']['lane_node_feats']
        lf[:, :, 0] = -lf[:, :, 0]  # x-coord
        lf[:, :, 2] = -lf[:, :, 2]  # yaw
        data['inputs']['map_representation']['lane_node_feats'] = lf

        # Flip surrounding agents
        vehicles = data['inputs']['surrounding_agent_representation']['vehicles']
        vehicles[:, :, 0] = -vehicles[:, :, 0]  # x-coord
        vehicles[:, :, 4] = -vehicles[:, :, 4]  # yaw-rate
        data['inputs']['surrounding_agent_representation']['vehicles'] = vehicles

        peds = data['inputs']['surrounding_agent_representation']['pedestrians']
        peds[:, :, 0] = -peds[:, :, 0]  # x-coord
        peds[:, :, 4] = -peds[:, :, 4]  # yaw-rate
        data['inputs']['surrounding_agent_representation']['pedestrians'] = peds

        # Flip groud truth trajectory
        fut = data['ground_truth']['traj']
        fut[:, 0] = -fut[:, 0]  # x-coord
        data['ground_truth']['traj'] = fut

        return data


# -*- coding: utf-8 -*-
"""
Created on Thu Oct 25 22:20:16 2018

@author: DELL
"""
import numpy as np
import cv2

img = cv2.imread("/Users/sarthakdandriyal/Desktop/ny.jpeg")
#img = cv2.imread("C:\\Users\\DELL\\Desktop\\proimg\\standard_test_images\\messi.jpg")

mask = np.zeros(img.shape[:2],np.uint8)

bgdModel = np.zeros((1,65),np.float64)
fgdModel = np.zeros((1,65),np.float64)

rect = (0,0,250,250)
#rect =(50,50,200,200)
cv2.grabCut(img,mask,rect,bgdModel,fgdModel,5,cv2.GC_INIT_WITH_RECT)
mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
img = img*mask2[:,:,np.newaxis]

cv2.imshow("Image", img)
cv2.waitKey(0)

#plt.imshow(img)
#plt.colorbar()
#plt.show()


from bibliopixel.animation.matrix import Matrix
from bibliopixel.util import log
import numpy as np
try:
    import cv2
except ImportError:
    log.error('Could not import cv2 library')
import os

grab = None

if os.name == 'nt':
    try:
        from desktopmagic.screengrab_win32 import getRectAsImage, getScreenAsImage
        log.debug("Using desktopmagic module")

        def nt_grab(bbox=None):
            if bbox is None:
                img = getScreenAsImage()
            else:
                img = getRectAsImage(bbox)
            return img

        grab = nt_grab
    except Exception:
        pass

if grab is None:
    try:
        from mss.linux import MSS as mss
        from PIL import Image

        sct = mss()
        monitor = sct.monitors[0]

        def mss_grab(bbox):
            sct_img = sct.grab(monitor)
            img = Image.frombytes('RGBA', sct_img.size, bytes(sct_img.raw), 'raw', 'BGRA').crop(bbox)
            img = img.convert('RGB')

            return img

        grab = mss_grab
        log.debug('Using mss module')
    except Exception:
        try:
            from PIL import ImageGrab
            grab = ImageGrab.grab
            log.debug("Using PIL ImageGrab module")
        except Exception:
            try:
                import pyscreenshot as ImageGrab
                grab = ImageGrab.grab
                log.debug("Using pyscreenshot module")
            except Exception:
                log.error("Unable to find any available screenshot option.")
                grab = None

class ScreenGrab(Matrix):

    def __init__(self, layout, bbox=(300, 300, 332, 332), mirror=False,
                 offset=0.0, crop=True, **kwds):
        super().__init__(layout, **kwds)

        if not sum(bbox):
            bbox = None
        self.bbox = bbox
        self.crop = crop
        self.mirror = mirror

        self.image = frame = self._capFrame()

        self._iw = frame.shape[1]
        self._ih = frame.shape[0]

        self.width = self.width
        self.height = self.height

        # self._scale = (self.height*1.0//self._ih)
        self._cropY = 0
        self._cropX = 0

        xoffset = yoffset = offset
        if xoffset > 1.0:
            xoffset = 1.0
        elif xoffset < -1.0:
            xoffset = -1.0
        if yoffset > 1.0:
            yoffset = 1.0
        elif yoffset < -1.0:
            yoffset = -1.0
        xoffset += 1.0
        yoffset += 1.0

        if self.height >= self.width:
            self._cropX = (self._iw - int(self.width / (self.height / float(self._ih)))) // 2
            if self._ih >= self._iw:
                scale = (self.height * 1.0) // self._ih
            else:
                scale = (self.width * 1.0) // self._iw
        else:
            self._cropY = (self._ih - int(self.height / (self.width / float(self._iw)))) // 2
            if self._ih >= self._iw:
                scale = (self.width * 1.0) // self._iw
            else:
                scale = (self.height * 1.0) // self._ih

        scaleW = int(self.width / scale)
        scaleH = int(self.height / scale)

        padTB = (scaleH - self._ih) // 2
        padLR = (scaleW - self._iw) // 2

        padYoff = int(round(padTB * yoffset)) - padTB
        padXoff = int(round(padLR * xoffset)) - padLR

        self._pad = (padTB + padYoff, padTB - padYoff,
                     padLR + padXoff, padLR - padXoff)

        self.xoff = int(round(self._cropX * xoffset)) - self._cropX
        self.yoff = int(round(self._cropY * yoffset)) - self._cropY

    def _capFrame(self):
        img = grab(self.bbox)
        return np.array(img)

    def step(self, amt=1):
        image = self._capFrame()

        if self.crop:
            image = image[self._cropY + self.yoff:self._ih - self._cropY +
                          self.yoff, self._cropX + self.xoff:self._iw - self._cropX + self.xoff]
        else:
            t, b, l, r = self._pad
            image = cv2.copyMakeBorder(
                image, t, b, l, r, cv2.BORDER_CONSTANT, value=[0, 0, 0])

        resized = cv2.resize(image, (self.width, self.height),
                             interpolation=cv2.INTER_LINEAR)
        if self.mirror:
            resized = cv2.flip(resized, 1)

        for y in range(self.height):
            for x in range(self.width):
                self.layout.set(x, y, tuple(resized[y, x][0:3]))


# coding: utf-8

from __future__ import division, unicode_literals
from functools import reduce

"""
This module implements representations of slabs and surfaces, as well as
algorithms for generating them.
"""

__author__ = "Richard Tran, Wenhao Sun, Zihan Xu, Shyue Ping Ong"
__copyright__ = "Copyright 2014, The Materials Virtual Lab"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "ongsp@ucsd.edu"
__date__ = "6/10/14"

from fractions import gcd
import math
import itertools
import logging

import numpy as np
from scipy.spatial.distance import squareform
from scipy.cluster.hierarchy import linkage, fcluster

from monty.fractions import lcm

from pymatgen.core.periodic_table import get_el_sp
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice
from pymatgen.core.sites import PeriodicSite

from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.util.coord_utils import in_coord_list
from pymatgen.analysis.structure_matcher import StructureMatcher

logger = logging.getLogger(__name__)

class Slab(Structure):
    """
    Subclass of Structure representing a Slab. Implements additional
    attributes pertaining to slabs, but the init method does not
    actually implement any algorithm that creates a slab. This is a
    DUMMY class who's init method only holds information about the
    slab. Also has additional methods that returns other information
    about a slab such as the surface area, normal, and atom adsorption.

    Note that all Slabs have the surface normal oriented in the c-direction.
    This means the lattice vectors a and b are in the surface plane and the c
    vector is out of the surface plane (though not necessary perpendicular to
    the surface.)

    .. attribute:: miller_index

        Miller index of plane parallel to surface.

    .. attribute:: scale_factor

        Final computed scale factor that brings the parent cell to the
        surface cell.

    .. attribute:: shift

        The shift value in Angstrom that indicates how much this
        slab has been shifted.
    """

    def __init__(self, lattice, species, coords, miller_index,
                 oriented_unit_cell, shift, scale_factor,
                 validate_proximity=False, to_unit_cell=False,
                 coords_are_cartesian=False, site_properties=None, energy=None):
        """
        Makes a Slab structure, a structure object with additional information
        and methods pertaining to slabs.

        Args:
            lattice (Lattice/3x3 array): The lattice, either as a
                :class:`pymatgen.core.lattice.Lattice` or
                simply as any 2D array. Each row should correspond to a lattice
                vector. E.g., [[10,0,0], [20,10,0], [0,0,30]] specifies a
                lattice with lattice vectors [10,0,0], [20,10,0] and [0,0,30].
            species ([Specie]): Sequence of species on each site. Can take in
                flexible input, including:

                i.  A sequence of element / specie specified either as string
                    symbols, e.g. ["Li", "Fe2+", "P", ...] or atomic numbers,
                    e.g., (3, 56, ...) or actual Element or Specie objects.

                ii. List of dict of elements/species and occupancies, e.g.,
                    [{"Fe" : 0.5, "Mn":0.5}, ...]. This allows the setup of
                    disordered structures.
            coords (Nx3 array): list of fractional/cartesian coordinates of
                each species.
            miller_index ([h, k, l]): Miller index of plane parallel to
                surface. Note that this is referenced to the input structure. If
                you need this to be based on the conventional cell,
                you should supply the conventional structure.
            oriented_unit_cell (Structure): The oriented_unit_cell from which
                this Slab is created (by scaling in the c-direction).
            shift (float): The shift in the c-direction applied to get the
                termination.
            scale_factor (array): scale_factor Final computed scale factor
                that brings the parent cell to the surface cell.
            validate_proximity (bool): Whether to check if there are sites
                that are less than 0.01 Ang apart. Defaults to False.
            coords_are_cartesian (bool): Set to True if you are providing
                coordinates in cartesian coordinates. Defaults to False.
            site_properties (dict): Properties associated with the sites as a
                dict of sequences, e.g., {"magmom":[5,5,5,5]}. The sequences
                have to be the same length as the atomic species and
                fractional_coords. Defaults to None for no properties.
            energy (float): A value for the energy.
        """
        self.oriented_unit_cell = oriented_unit_cell
        self.miller_index = tuple(miller_index)
        self.shift = shift
        self.scale_factor = scale_factor
        self.energy = energy
        super(Slab, self).__init__(
            lattice, species, coords, validate_proximity=validate_proximity,
            to_unit_cell=to_unit_cell,
            coords_are_cartesian=coords_are_cartesian,
            site_properties=site_properties)

    def get_sorted_structure(self, key=None, reverse=False):
        """
        Get a sorted copy of the structure. The parameters have the same
        meaning as in list.sort. By default, sites are sorted by the
        electronegativity of the species. Note that Slab has to override this
        because of the different __init__ args.

        Args:
            key: Specifies a function of one argument that is used to extract
                a comparison key from each list element: key=str.lower. The
                default value is None (compare the elements directly).
            reverse (bool): If set to True, then the list elements are sorted
                as if each comparison were reversed.
        """
        sites = sorted(self, key=key, reverse=reverse)
        s = Structure.from_sites(sites)
        return Slab(s.lattice, s.species_and_occu, s.frac_coords,
                    self.miller_index, self.oriented_unit_cell, self.shift,
                    self.scale_factor, site_properties=s.site_properties)

    @property
    def normal(self):
        """
        Calculates the surface normal vector of the slab
        """
        normal = np.cross(self.lattice.matrix[0], self.lattice.matrix[1])
        normal /= np.linalg.norm(normal)
        return normal

    @property
    def surface_area(self):
        """
        Calculates the surface area of the slab
        """
        m = self.lattice.matrix
        return np.linalg.norm(np.cross(m[0], m[1]))

    def add_adsorbate_atom(self, indices, specie, distance):
        """
        Gets the structure of single atom adsorption.
        slab structure from the Slab class(in [0, 0, 1])

        Args:
            indices ([int]): Indices of sites on which to put the absorbate.
                Absorbed atom will be displaced relative to the center of
                these sites.
            specie (Specie/Element/str): adsorbed atom species
            distance (float): between centers of the adsorbed atom and the
                given site in Angstroms.
        """
        #Let's do the work in cartesian coords
        center = np.sum([self[i].coords for i in indices], axis=0) / len(
            indices)

        coords = center + self.normal * distance / np.linalg.norm(self.normal)

        self.append(specie, coords, coords_are_cartesian=True)

    def __str__(self):
        comp = self.composition
        outs = [
            "Slab Summary (%s)" % comp.formula,
            "Reduced Formula: %s" % comp.reduced_formula,
            "Miller index: %s" % (self.miller_index, ),
            "Shift: %.4f, Scale Factor: %s" % (self.shift,
                                               self.scale_factor.__str__())]
        to_s = lambda x: "%0.6f" % x
        outs.append("abc   : " + " ".join([to_s(i).rjust(10)
                                           for i in self.lattice.abc]))
        outs.append("angles: " + " ".join([to_s(i).rjust(10)
                                           for i in self.lattice.angles]))
        outs.append("Sites ({i})".format(i=len(self)))
        for i, site in enumerate(self):
            outs.append(" ".join([str(i + 1), site.species_string,
                                  " ".join([to_s(j).rjust(12)
                                            for j in site.frac_coords])]))
        return "\n".join(outs)

    def as_dict(self):
        d = super(Slab, self).as_dict()
        d["@module"] = self.__class__.__module__
        d["@class"] = self.__class__.__name__
        d["oriented_unit_cell"] = self.oriented_unit_cell.as_dict()
        d["miller_index"] = self.miller_index
        d["shift"] = self.shift
        d["scale_factor"] = self.scale_factor
        d["energy"] = self.energy
        return d

    @classmethod
    def from_dict(cls, d):
        lattice = Lattice.from_dict(d["lattice"])
        sites = [PeriodicSite.from_dict(sd, lattice) for sd in d["sites"]]
        s = Structure.from_sites(sites)

        return Slab(
            lattice=lattice,
            species=s.species_and_occu, coords=s.frac_coords,
            miller_index=d["miller_index"],
            oriented_unit_cell=Structure.from_dict(d["oriented_unit_cell"]),
            shift=d["shift"], scale_factor=d["scale_factor"],
            site_properties=s.site_properties, energy=d["energy"]
        )

class SlabGenerator(object):

    """
    This class generates different slabs using shift values determined by where a
    unique termination can be found along with other criterias such as where a
    termination doesn't break a polyhedral bond. The shift value then indicates
    where the slab layer will begin and terminate in the slab-vacuum system.

    .. attribute:: oriented_unit_cell

        A unit cell of the parent structure with the miller
        index of plane parallel to surface

    .. attribute:: parent

        Parent structure from which Slab was derived.

    .. attribute:: lll_reduce

        Whether or not the slabs will be orthogonalized

    .. attribute:: center_slab

        Whether or not the slabs will be centered between
        the vacuum layer

    .. attribute:: slab_scale_factor

        Final computed scale factor that brings the parent cell to the
        surface cell.

    .. attribute:: miller_index

        Miller index of plane parallel to surface.

    .. attribute:: min_slab_size

        Minimum size in angstroms of layers containing atoms

    .. attribute:: min_vac_size

        Minimize size in angstroms of layers containing vacuum

    """

    def __init__(self, initial_structure, miller_index, min_slab_size,
                 min_vacuum_size, lll_reduce=False, center_slab=False,
                 primitive=True):
        """
        Calculates the slab scale factor and uses it to generate a unit cell
        of the initial structure that has been oriented by its miller index.
        Also stores the initial information needed later on to generate a slab.

        Args:
            initial_structure (Structure): Initial input structure.
            miller_index ([h, k, l]): Miller index of plane parallel to
                surface. Note that this is referenced to the input structure. If
                you need this to be based on the conventional cell,
                you should supply the conventional structure.
            min_slab_size (float): In Angstroms
            min_vac_size (float): In Angstroms
            lll_reduce (bool): Whether to perform an LLL reduction on the
                eventual structure.
            center_slab (bool): Whether to center the slab in the cell with
                equal vacuum spacing from the top and bottom.
            primitive (bool): Whether to reduce any generated slabs to a
                primitive cell (this does **not** mean the slab is generated
                from a primitive cell, it simply means that after slab
                generation, we attempt to find shorter lattice vectors,
                which lead to less surface area and smaller cells).
        """
        latt = initial_structure.lattice
        d = abs(reduce(gcd, miller_index))
        miller_index = tuple([int(i / d) for i in miller_index])
        #Calculate the surface normal using the reciprocal lattice vector.
        recp = latt.reciprocal_lattice_crystallographic
        normal = recp.get_cartesian_coords(miller_index)
        normal /= np.linalg.norm(normal)

        slab_scale_factor = []
        non_orth_ind = []
        eye = np.eye(3, dtype=np.int)
        for i, j in enumerate(miller_index):
            if j == 0:
                # Lattice vector is perpendicular to surface normal, i.e.,
                # in plane of surface. We will simply choose this lattice
                # vector as one of the basis vectors.
                slab_scale_factor.append(eye[i])
            else:
                #Calculate projection of lattice vector onto surface normal.
                d = abs(np.dot(normal, latt.matrix[i])) / latt.abc[i]
                non_orth_ind.append((i, d))

        # We want the vector that has maximum magnitude in the
        # direction of the surface normal as the c-direction.
        # Results in a more "orthogonal" unit cell.
        c_index, dist = max(non_orth_ind, key=lambda t: t[1])

        if len(non_orth_ind) > 1:
            lcm_miller = lcm(*[miller_index[i] for i, d in non_orth_ind])
            for (i, di), (j, dj) in itertools.combinations(non_orth_ind, 2):
                l = [0, 0, 0]
                l[i] = -int(round(lcm_miller / miller_index[i]))
                l[j] = int(round(lcm_miller / miller_index[j]))
                slab_scale_factor.append(l)
                if len(slab_scale_factor) == 2:
                    break

        slab_scale_factor.append(eye[c_index])

        slab_scale_factor = np.array(slab_scale_factor)

        # Let's make sure we have a left-handed crystallographic system
        if np.linalg.det(slab_scale_factor) < 0:
            slab_scale_factor *= -1

        single = initial_structure.copy()
        single.make_supercell(slab_scale_factor)

        self.oriented_unit_cell = Structure.from_sites(single,
                                                       to_unit_cell=True)
        self.parent = initial_structure
        self.lll_reduce = lll_reduce
        self.center_slab = center_slab
        self.slab_scale_factor = slab_scale_factor
        self.miller_index = miller_index
        self.min_vac_size = min_vacuum_size
        self.min_slab_size = min_slab_size
        self.primitive = primitive
        self._normal = normal
        a, b, c = self.oriented_unit_cell.lattice.matrix
        self._proj_height = abs(np.dot(normal, c))

    def get_slab(self, shift=0, tol=0.1, energy=None):
        """
        This method takes in shift value for the c lattice direction and
        generates a slab based on the given shift. You should rarely use this
        method. Instead, it is used by other generation algorithms to obtain
        all slabs.

        Arg:
            shift (float): A shift value in Angstrom that determines how much a
                slab should be shifted.
            tol (float): Tolerance to determine primitive cell.
            energy (float): An energy to assign to the slab.

        Returns:
            (Slab) A Slab object with a particular shifted oriented unit cell.
        """

        h = self._proj_height
        nlayers_slab = int(math.ceil(self.min_slab_size / h))
        nlayers_vac = int(math.ceil(self.min_vac_size / h))
        nlayers = nlayers_slab + nlayers_vac

        species = self.oriented_unit_cell.species_and_occu
        props = self.oriented_unit_cell.site_properties
        props = {k: v * nlayers_slab for k, v in props.items()}
        frac_coords = self.oriented_unit_cell.frac_coords
        frac_coords = np.array(frac_coords) +\
                      np.array([0, 0, -shift])[None, :]
        frac_coords = frac_coords - np.floor(frac_coords)
        a, b, c = self.oriented_unit_cell.lattice.matrix
        new_lattice = [a, b, nlayers * c]
        frac_coords[:, 2] = frac_coords[:, 2] / nlayers
        all_coords = []
        for i in range(nlayers_slab):
            fcoords = frac_coords.copy()
            fcoords[:, 2] += i / nlayers
            all_coords.extend(fcoords)

        slab = Structure(new_lattice, species * nlayers_slab, all_coords,
                         site_properties=props)

        scale_factor = self.slab_scale_factor
        # Whether or not to orthogonalize the structure
        if self.lll_reduce:
            lll_slab = slab.copy(sanitize=True)
            mapping = lll_slab.lattice.find_mapping(slab.lattice)
            scale_factor = np.dot(mapping[2], scale_factor)
            slab = lll_slab

        # Whether or not to center the slab layer around the vacuum
        if self.center_slab:
            avg_c = np.average([c[2] for c in slab.frac_coords])
            slab.translate_sites(list(range(len(slab))), [0, 0, 0.5 - avg_c])

        if self.primitive:
            prim = slab.get_primitive_structure(tolerance=tol)
            if energy is not None:
                energy = prim.volume / slab.volume * energy
            slab = prim

        return Slab(slab.lattice, slab.species_and_occu,
                    slab.frac_coords, self.miller_index,
                    self.oriented_unit_cell, shift,
                    scale_factor, site_properties=slab.site_properties,
                    energy=energy)

    def _calculate_possible_shifts(self, tol=0.1):
        frac_coords = self.oriented_unit_cell.frac_coords
        n = len(frac_coords)

        # We cluster the sites according to the c coordinates. But we need to
        # take into account PBC. Let's compute a fractional c-coordinate
        # distance matrix that accounts for PBC.
        dist_matrix = np.zeros((n, n))
        h = self._proj_height
        # Projection of c lattice vector in
        # direction of surface normal.
        for i, j in itertools.combinations(list(range(n)), 2):
            if i != j:
                cdist = frac_coords[i][2] - frac_coords[j][2]
                cdist = abs(cdist - round(cdist)) * h
                dist_matrix[i, j] = cdist
                dist_matrix[j, i] = cdist

        condensed_m = squareform(dist_matrix)
        z = linkage(condensed_m)
        clusters = fcluster(z, tol, criterion="distance")

        #Generate dict of cluster# to c val - doesn't matter what the c is.
        c_loc = {c: frac_coords[i][2] for i, c in enumerate(clusters)}

        #Put all c into the unit cell.
        possible_c = [c - math.floor(c) for c in sorted(c_loc.values())]

        # Calculate the shifts
        nshifts = len(possible_c)
        shifts = []
        for i in range(nshifts):
            if i == nshifts - 1:
                # There is an additional shift between the first and last c
                # coordinate. But this needs special handling because of PBC.
                shift = (possible_c[0] + 1 + possible_c[i]) * 0.5
                if shift > 1:
                    shift -= 1
            else:
                shift = (possible_c[i] + possible_c[i + 1]) * 0.5
            shifts.append(shift - math.floor(shift))
        shifts = sorted(shifts)
        return shifts

    def get_slabs(self, bonds=None, tol=0.1, max_broken_bonds=0):
        """
        This method returns a list of slabs that are generated using the list of
        shift values from the method, _calculate_possible_shifts(). Before the
        shifts are used to create the slabs however, if the user decides to take
        into account whether or not a termination will break any polyhedral
        structure (bonds != None), this method will filter out any shift values
        that do so.

        Args:
            bonds ({(specie1, specie2): max_bond_dist}: bonds are
                specified as a dict of tuples: float of specie1, specie2
                and the max bonding distance. For example, PO4 groups may be
                defined as {("P", "O"): 3}.
            tol (float): Threshold parameter in fcluster in order to check
                if two atoms are lying on the same plane. Default thresh set
                to 0.1 Angstrom in the direction of the surface normal.
            max_broken_bonds (int): Maximum number of allowable broken bonds
                for the slab. Use this to limit # of slabs (some structures
                may have a lot of slabs). Defaults to zero, which means no
                defined bonds must be broken.

        Returns:
            ([Slab]) List of all possible terminations of a particular surface.
            Slabs are sorted by the # of bonds broken.
        """
        def get_c_ranges(site1, site2, bond_dist):
            lattice = site1.lattice
            f1 = site1.frac_coords
            c_ranges = []
            for dist, image in lattice.get_all_distance_and_image(
                    f1, site2.frac_coords):
                if dist < bond_dist:
                    f2 = site2.frac_coords + image
                    # Checks if the distance between the two species
                    # is less then the user input bond distance
                    min_c = f1[2]
                    max_c = f2[2]
                    c_range = sorted([min_c, max_c])
                    if c_range[1] > 1:
                        # Takes care of PBC when c coordinate of site
                        # goes beyond the upper boundary of the cell
                        c_ranges.append((c_range[0], 1))
                        c_ranges.append((0, c_range[1] - 1))
                    elif c_range[0] < 0:
                        # Takes care of PBC when c coordinate of site
                        # is below the lower boundary of the unit cell
                        c_ranges.append((0, c_range[1]))
                        c_ranges.append((c_range[0] + 1, 1))
                    else:
                        c_ranges.append(c_range)
            return c_ranges

        bond_c_ranges = []
        if bonds is not None:
            #Convert to species first
            bonds = {(get_el_sp(s1), get_el_sp(s2)): dist for (s1, s2), dist in
                     bonds.items()}
            for s1, s2 in itertools.combinations(self.oriented_unit_cell, 2):
                # Iterates through every possible pair of species in the
                # oriented unit cell
                all_sp = set(s1.species_and_occu.keys())
                all_sp.update(s2.species_and_occu.keys())
                for species, bond_dist in bonds.items():
                    # Checks if elements in species is in all_sp
                    if all_sp.issuperset(species):
                        bond_c_ranges.extend(get_c_ranges(s1, s2, bond_dist))

        slabs = []
        for shift in self._calculate_possible_shifts(tol=tol):
            bonds_broken = 0
            for r in bond_c_ranges:
                if r[0] <= shift <= r[1]:
                    bonds_broken += 1
            if bonds_broken <= max_broken_bonds:
                # For now, set the energy to be equal to no. of broken bonds
                # per unit cell.
                slab = self.get_slab(shift, tol=tol, energy=bonds_broken)
                slabs.append(slab)

        # Further filters out any surfaces made that might be the same
        m = StructureMatcher(ltol=tol, stol=tol, primitive_cell=False,
                             scale=False)
        slabs = [g[0] for g in m.group_structures(slabs)]
        return sorted(slabs, key=lambda s: s.energy)

def get_symmetrically_distinct_miller_indices(structure, max_index):
    """
    Returns all symmetrically distinct indices below a certain max-index for
    a given structure. Analysis is based on the symmetry of the reciprocal
    lattice of the structure.

    Args:
        structure (Structure): input structure.
        max_index (int): The maximum index. For example, a max_index of 1
            means that (100), (110), and (111) are returned for the cubic
            structure. All other indices are equivalent to one of these.
    """

    recp_lattice = structure.lattice.reciprocal_lattice_crystallographic
    # Need to make sure recp lattice is big enough, otherwise symmetry
    # determination will fail. We set the overall volume to 1.
    recp_lattice = recp_lattice.scale(1)
    recp = Structure(recp_lattice, ["H"], [[0, 0, 0]])

    # Creates a function that uses the symmetry operations in the
    # structure to find Miller indices that might give repetitive slabs
    analyzer = SpacegroupAnalyzer(recp, symprec=0.001)
    symm_ops = analyzer.get_symmetry_operations()
    unique_millers = []

    def is_already_analyzed(miller_index):
        for op in symm_ops:
            if in_coord_list(unique_millers, op.operate(miller_index)):
                return True
        return False

    r = list(range(-max_index, max_index + 1))
    r.reverse()
    for miller in itertools.product(r, r, r):
        if any([i != 0 for i in miller]):
            d = abs(reduce(gcd, miller))
            miller = tuple([int(i / d) for i in miller])
            if not is_already_analyzed(miller):
                unique_millers.append(miller)
    return unique_millers

def generate_all_slabs(structure, max_index, min_slab_size, min_vacuum_size,
                       bonds=None, tol=0.1, max_broken_bonds=0,
                       lll_reduce=False,
                       center_slab=False, primitive=True):
    """
    A function that finds all different slabs up to a certain miller index.
    Slabs oriented under certain Miller indices that are equivalent to other
    slabs in other Miller indices are filtered out using symmetry operations
    to get rid of any repetitive slabs. For example, under symmetry operations,
    CsCl has equivalent slabs in the (0,0,1), (0,1,0), and (1,0,0) direction.

    Args:
        structure (Structure): Initial input structure.
        max_index (int): The maximum Miller index to go up to.
        symprec (float): Tolerance for symmetry finding. Defaults to 1e-3,
            which is fairly strict and works well for properly refined
            structures with atoms in the proper symmetry coordinates. For
            structures with slight deviations from their proper atomic
            positions (e.g., structures relaxed with electronic structure
            codes), a looser tolerance of 0.1 (the value used in Materials
            Project) is often needed.
        lll_reduce (bool): Whether to perform an LLL reduction on the
            eventual structure.
        center_slab (bool): Whether to center the slab in the cell with
            equal vacuum spacing from the top and bottom.
        primitive (bool): Whether to reduce any generated slabs to a
            primitive cell (this does **not** mean the slab is generated
            from a primitive cell, it simply means that after slab
            generation, we attempt to find shorter lattice vectors,
            which lead to less surface area and smaller cells).
    """
    all_slabs = []
    for miller in get_symmetrically_distinct_miller_indices(structure,
                                                            max_index):
        gen = SlabGenerator(structure, miller, min_slab_size,
                            min_vacuum_size, lll_reduce=lll_reduce,
                            center_slab=center_slab, primitive=primitive)
        slabs = gen.get_slabs(bonds=bonds, tol=tol,
                              max_broken_bonds=max_broken_bonds)
        if len(slabs) > 0:
            logger.debug("%s has %d slabs... " % (miller, len(slabs)))
            all_slabs.extend(slabs)

    return all_slabs