import pygame
import random
import itertools
import report
import math

'''
Define constants
'''
SCREEN_WIDTH = 1000
SCREEN_HEIGHT = 500
SCREEN_COLOR = (0, 0, 0) # Black
LANE_WIDTH = 19
LANE_COLOR = (220, 220, 220) # Gray
LANE_NAME_COLOR = (255, 255, 255) # White

HORIZONTAL = 1
VERTICAL = 2
TO_LEFT = 3
TO_RIGHT= 4
TO_BOTTOM = 5
TO_TOP = 6

GO = 100
REDLIGHT = 101  # Red light
BLOCKED = 102   # Preceding car is near

CAR_WIDTH = 10
CAR_LENGTH = 20
CAR_COLOR = (50, 200, 50) # Greenish
CAR_COLOR_ACCIDENT = (200, 50, 50) # Redish
CAR_COLOR_VAR = 50
CAR_SPEED = 10
CAR_SPEED_VAR = 3
CAR_SAFE_DISTANCE = CAR_LENGTH * 1.5
CAR_CHANGE_LANE_RATE_BLOCKED = 0.2  # Rate of chaining lanes when blocked

'''
Defind a Road object
'''
class Road():
    '''
    Constructor of a Road object
    orientation: HORIZONTAL or VERTICAL
    num_lanes: # of to-left and to-right lanes for a HORIZONTAL road
                            # of to-bottom and to-top lanes for a VERTICAL road
    position: top y coordinate for a horizontal road
                     left x coordinate for a vertical road
    '''
    def __init__(self, pygame, name, font, position, orientation, num_lanes):
        self.pygame = pygame
        self.name = font.render(name, True, LANE_NAME_COLOR)
        self.name_str = name
        self.position = position
        self.intersections = []
        self.lanes_for_new_cars = []

        if orientation != HORIZONTAL and orientation != VERTICAL:
            raise ValueError('Illegal orientation is used')
        self.orientation = orientation

        # Sanity check for num_lanes
        if len(num_lanes) != 2:
            raise ValueError('num_lanes requires two integers')
        if num_lanes[0]<0 or num_lanes[1]<0:
            raise ValueError('num_lanes requires two integers >=0')
        if (num_lanes[0]==0 and num_lanes[1]==0):
            raise ValueError('At least one element of num_lanes must be >0')

        # Add lanes
        self.lanes = []
        if self.orientation == HORIZONTAL:
            for idx in range(num_lanes[0]):
                self.lanes.append(Lane(pygame, TO_LEFT, position + (self.name.get_height() + 2) + idx * (LANE_WIDTH+1), self))
            for idx in range(num_lanes[1]):
                self.lanes.append(Lane(pygame, TO_RIGHT, position + (self.name.get_height() + 2) + (idx+num_lanes[0]) * (LANE_WIDTH+1), self))
        elif self.orientation == VERTICAL:
            for idx in range(num_lanes[0]):
                self.lanes.append(Lane(pygame, TO_BOTTOM, position + (self.name.get_width() + 2) + idx * (LANE_WIDTH+1), self))
            for idx in range(num_lanes[1]):
                self.lanes.append(Lane(pygame, TO_TOP, position + (self.name.get_width() + 2) + (idx+num_lanes[1]) * (LANE_WIDTH+1), self))

        configure_lanes_before_after(self.lanes)

    def add_newCar(self):
        if len(self.lanes_for_new_cars) == 1:
            self.lanes_for_new_cars[0].add_newCar()
        elif len(self.lanes_for_new_cars) > 1:
            lane = self.lanes_for_new_cars[random.randrange(0, len(self.lanes_for_new_cars))]
            lane.add_newCar()

    def paint_on(self, screen):
        for lane in self.lanes:
            lane.paint_on(screen)
        if self.orientation == HORIZONTAL:
            screen.blit(self.name, (0, self.position))
        elif self.orientation == VERTICAL:
            screen.blit(self.name, (self.position, 20))

    def paint_cars_on(self, screen):
        for lane in self.lanes:
            lane.paint_cars_on(screen)

    def move(self, batch):
        for lane in self.lanes:
            lane.move(batch)

    def find_lanes_on_mouse_pos(self, x, y):
        lanes = []
        for lane in self.lanes:
            if lane.include_pos(x, y):
                lanes.append(lane)
        return lanes

    def on_the_same_road_with(self, road):
        if isinstance(road, Road):
            return self == road
        elif isinstance(road, Intersection):
            for r in road.roads:
                if self == r:
                    return True
            return False
        else:
            return False

    def intersect(self, road):
        if self.orientation == HORIZONTAL and road.orientation == VERTICAL:
            return self, road
        elif self.orientation == VERTICAL and road.orientation == HORIZONTAL:
            return road, self
        return None

    def overlap(self, road):
        if self.orientation == HORIZONTAL and road.orientation == HORIZONTAL:
            max_top = max(self.position, road.position)
            min_bottom = min(self.lanes[-1].rect.bottom, road.lanes[-1].rect.bottom)
            if max_top <= min_bottom:
                return True
        elif self.orientation == VERTICAL and road.orientation == VERTICAL:
            max_left = max(self.position, road.position)
            min_right = min(self.lanes[-1].rect.right, road.lanes[-1].rect.right)
            if max_left <= min_right:
                return True
        return False

'''
For each pair of roads with same orientation,
                    check to see whether they overlap
'''
def find_overlaps(roads):
    comb = itertools.combinations(roads, 2)
    for pair in comb:
        if pair[0].overlap(pair[1]):
            raise ValueError(pair[0].name_str + ' and ' + pair[1].name_str + ' overlap')

'''
Find lanes where new cars can be added
'''
def find_lanes_for_new_cars(roads):
    for road in roads:
        if isinstance(road, Intersection):
            continue

        road.lanes_for_new_cars.clear()
        for lane in road.lanes:
            if (lane.direction == TO_LEFT and lane.rect.right == SCREEN_WIDTH) or\
                (lane.direction == TO_RIGHT and lane.rect.left == 0) or\
                (lane.direction == TO_BOTTOM and lane.rect.top == 0) or\
                (lane.direction == TO_TOP and lane.rect.bottom == SCREEN_HEIGHT):
                road.lanes_for_new_cars.append(lane)

'''
Define a Lane object
'''
class Lane():
    '''
    Constructor of a Lane object
    position: top y coordinate for a horizontal lane
                     left x coordinate for a vertical lane
    '''
    def __init__(self, pygame, direction, position, road):
        self.pygame = pygame

        # Position the lane and get its Rectangle object
        self.direction = direction
        self.position = position
        self.center = position + LANE_WIDTH/2
        if self.direction == TO_LEFT or self.direction == TO_RIGHT:
            self.surf = pygame.Surface((SCREEN_WIDTH, LANE_WIDTH))  # X/Y size
            self.rect = self.surf.get_rect(center=(SCREEN_WIDTH/2, self.center))
        elif self.direction == TO_BOTTOM or self.direction == TO_TOP:
            self.surf = pygame.Surface((LANE_WIDTH, SCREEN_HEIGHT))  # X/Y size
            self.rect = self.surf.get_rect(center=(self.center, SCREEN_HEIGHT/2))
        self.surf.fill(LANE_COLOR)

        self.road = road
        self.cars = []
        self.before = None
        self.after = None
        self.next = []      # connected lanes
        self.trafficLight = GO

        self.blocking_lanes = []                 # Other lanes that cross (thus possibly block) this lane

    def update_size(self, left, top, width, height):
        self.rect.update(left, top, width, height)
        self.surf = pygame.transform.scale(self.surf, (width, height))

    def add_newCar(self):
        if self.direction == TO_LEFT and self.rect.right == SCREEN_WIDTH:
            if (len(self.cars) == 0) or (self.cars[-1].rect.right < self.rect.right - CAR_SAFE_DISTANCE):
                self.cars.append(Car(self.pygame, self.road, self, self.rect.right - CAR_LENGTH/2, self.center))
        elif self.direction == TO_RIGHT and self.rect.left == 0:
            if (len(self.cars) == 0) or (self.cars[-1].rect.left > self.rect.left + CAR_SAFE_DISTANCE):
                self.cars.append(Car(self.pygame, self.road, self, self.rect.left + CAR_LENGTH/2, self.center))
        elif self.direction == TO_BOTTOM and self.rect.top == 0:
            if (len(self.cars) == 0) or (self.cars[-1].rect.top > self.rect.top + CAR_SAFE_DISTANCE):
                self.cars.append(Car(self.pygame, self.road, self, self.center, self.rect.top + CAR_LENGTH/2))
        elif self.direction == TO_TOP and self.rect.bottom == SCREEN_HEIGHT:
            if (len(self.cars) == 0) or (self.cars[-1].rect.bottom < self.rect.bottom - CAR_SAFE_DISTANCE):
                self.cars.append(Car(self.pygame, self.road, self, self.center, self.rect.bottom - CAR_LENGTH/2))

    def paint_on(self, screen):
        screen.blit(self.surf, self.rect)

    def paint_cars_on(self, screen):
        for car in self.cars:
            # Adjust each car's center to the lane's center, as the car might have changed lanes
            if self.direction == TO_LEFT or self.direction == TO_RIGHT:
                car.rect.centery = self.center
            elif self.direction == TO_BOTTOM or self.direction == TO_TOP:
                car.rect.centerx = self.center
            car.paint_on(screen)

    def move(self, batch):
        # Initialize position of preceding car for keeping a safe distance between consecutive cars
        if self.direction == TO_LEFT:
            self.x_preceding_car = self.rect.left - (CAR_SAFE_DISTANCE * 2)
        elif self.direction == TO_RIGHT:
            self.x_preceding_car = self.rect.right + (CAR_SAFE_DISTANCE * 2)
        elif self.direction == TO_BOTTOM:
            self.y_preceding_car = self.rect.bottom + (CAR_SAFE_DISTANCE * 2)
        elif self.direction == TO_TOP:
            self.y_preceding_car = self.rect.top - (CAR_SAFE_DISTANCE * 2)

        self.status_preceding_car = GO

        # Move each car on the lane
        self.cars = [car for car in self.cars if car.move(batch)]        # Only cars visible on the screen remain in the list

    '''
    Return (True, idx) if car can be inserted into this lane's cars[idx]
    Otherwise, return (False, None)
    '''
    def can_change_lane(self, car, add):
        safe_distance = CAR_SAFE_DISTANCE * 2

        if self.direction == TO_LEFT:
            # Check to see if cars on this lane blocks entrance
            id_nearest_car = self.find_nearest_car_to_left(car, 0, len(self.cars)-1)
            if ( (id_nearest_car >= 0) and (car.rect.left < self.cars[id_nearest_car].rect.right + safe_distance) ) or\
               ( (id_nearest_car+1 < len(self.cars)) and (self.cars[id_nearest_car+1].rect.left - safe_distance < car.rect.right)):
                return False, None

            # Check to see if cars on lanes that cross this lane blocks entrance
            for blocking_lane in self.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.include_car(c) and not car.outside_safe_distance_from(c, safe_distance, self.direction):
                        return False, None

            # Check to see if cars on this lane's next lane blocks entrance
            if car.rect.left - car.speed <= self.rect.left + safe_distance:
                for next_lane in self.next:
                    if next_lane.trafficLight == REDLIGHT and car.rect.left < next_lane.rect.right + safe_distance:
                        return False, None
                    if len(next_lane.cars) > 0 and car.rect.left < next_lane.cars[-1].rect.right + safe_distance:
                        return False, None
                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c) and car.rect.left < c.rect.right + safe_distance:
                                return False, None

            return True, id_nearest_car+1

        elif self.direction == TO_RIGHT:
            # Check to see if cars on this lane blocks entrance
            id_nearest_car = self.find_nearest_car_to_right(car, 0, len(self.cars)-1)
            if ( (id_nearest_car >= 0) and (self.cars[id_nearest_car].rect.left - safe_distance < car.rect.right ) ) or\
               ( (id_nearest_car+1 < len(self.cars)) and (car.rect.left < self.cars[id_nearest_car+1].rect.right + safe_distance)):
                return False, None

            # Check to see if cars on lanes that cross this lane blocks entrance
            for blocking_lane in self.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.include_car(c) and not car.outside_safe_distance_from(c, safe_distance, self.direction):
                        return False, None

            # Check to see if cars on this lane's next lane blocks entrance
            if self.rect.right - safe_distance <= car.rect.right + car.speed:
                for next_lane in self.next:
                    if next_lane.trafficLight == REDLIGHT and next_lane.rect.left - safe_distance < car.rect.right:
                        return False, None
                    if len(next_lane.cars) > 0 and next_lane.cars[-1].rect.left - safe_distance < car.rect.right:
                        return False, None
                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c) and c.rect.left - safe_distance < car.rect.right:
                                return False, None

            return True, id_nearest_car+1

        elif self.direction == TO_BOTTOM:
            # Check to see if cars on this lane blocks entrance
            id_nearest_car = self.find_nearest_car_to_bottom(car, 0, len(self.cars)-1)
            if ( (id_nearest_car >= 0) and (self.cars[id_nearest_car].rect.top - safe_distance < car.rect.bottom ) ) or\
               ( (id_nearest_car+1 < len(self.cars)) and (car.rect.top < self.cars[id_nearest_car+1].rect.bottom + safe_distance)):
                return False, None

            # Check to see if cars on lanes that cross this lane blocks entrance
            for blocking_lane in self.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.include_car(c) and not car.outside_safe_distance_from(c, safe_distance, self.direction):
                        return False, None

            # Check to see if cars on this lane's next lane blocks entrance
            if self.rect.bottom - safe_distance <= car.rect.bottom + car.speed:
                for next_lane in self.next:
                    if next_lane.trafficLight == REDLIGHT and next_lane.rect.top - safe_distance < car.rect.bottom:
                        return False, None
                    if len(next_lane.cars) > 0 and next_lane.cars[-1].rect.top - safe_distance < car.rect.bottom:
                        return False, None
                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c) and c.rect.top - safe_distance < car.rect.bottom:
                                return False, None

            return True, id_nearest_car+1

        elif self.direction == TO_TOP:
            # Check to see if cars on this lane blocks entrance
            id_nearest_car = self.find_nearest_car_to_top(car, 0, len(self.cars)-1)
            if ( (id_nearest_car >= 0) and (car.rect.top < self.cars[id_nearest_car].rect.bottom + safe_distance) ) or\
               ( (id_nearest_car+1 < len(self.cars)) and (self.cars[id_nearest_car+1].rect.top - safe_distance < car.rect.bottom)):
                return False, None

            # Check to see if cars on lanes that cross this lane blocks entrance
            for blocking_lane in self.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.include_car(c) and not car.outside_safe_distance_from(c, safe_distance, self.direction):
                        return False, None

            # Check to see if cars on this lane's next lane blocks entrance
            if car.rect.top - car.speed <= self.rect.top + safe_distance:
                for next_lane in self.next:
                    if next_lane.trafficLight == REDLIGHT and car.rect.top < next_lane.rect.bottom + safe_distance:
                        return False, None
                    if len(next_lane.cars) > 0 and car.rect.top < next_lane.cars[-1].rect.bottom + safe_distance:
                        return False, None
                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c) and car.rect.top < c.rect.bottom + safe_distance:
                                return False, None

            return True, id_nearest_car+1

    def find_nearest_car_to_left(self, car, id_left, id_right):         # Perform binary search to find the nearest car
        if (id_left > id_right):
            return id_right   # If a car with the same position is not found, return the car to the nearest front

        id_mid = int((id_left+id_right)/2)
        if (self.cars[id_mid].rect.left == car.rect.left):
            return id_mid
        elif (self.cars[id_mid].rect.left < car.rect.left):
            return self.find_nearest_car_to_left(car, id_mid+1, id_right)
        else:
            return self.find_nearest_car_to_left(car, id_left, id_mid-1)

    def find_nearest_car_to_right(self, car, id_left, id_right):         # Perform binary search to find the nearest car
        if (id_left > id_right):
            return id_right   # If a car with the same position is not found, return the car to the nearest front

        id_mid = int((id_left+id_right)/2)
        if (self.cars[id_mid].rect.left == car.rect.left):
            return id_mid
        elif (self.cars[id_mid].rect.left < car.rect.left):
            return self.find_nearest_car_to_right(car, id_left, id_mid-1)
        else:
            return self.find_nearest_car_to_right(car, id_mid+1, id_right)

    def find_nearest_car_to_bottom(self, car, id_top, id_bottom):         # Perform binary search to find the nearest car
        if (id_top > id_bottom):
            return id_bottom   # If a car with the same position is not found, return the car to the nearest front

        id_mid = int((id_top+id_bottom)/2)
        if (self.cars[id_mid].rect.top == car.rect.top):
            return id_mid
        elif (self.cars[id_mid].rect.top < car.rect.top):
            return self.find_nearest_car_to_bottom(car, id_top, id_mid-1)
        else:
            return self.find_nearest_car_to_bottom(car, id_mid+1, id_bottom)

    def find_nearest_car_to_top(self, car, id_top, id_bottom):         # Perform binary search to find the nearest car
        if (id_top > id_bottom):
            return id_bottom   # If a car with the same position is not found, return the car to the nearest front

        id_mid = int((id_top+id_bottom)/2)
        if (self.cars[id_mid].rect.top == car.rect.top):
            return id_mid
        elif (self.cars[id_mid].rect.top < car.rect.top):
            return self.find_nearest_car_to_left(car, id_mid+1, id_bottom)
        else:
            return self.find_nearest_car_to_left(car, id_top, id_mid-1)

    def include_pos(self, x, y):
        if (self.rect.left <= x) and (x <= self.rect.right) and\
               (self.rect.top <= y) and (y <= self.rect.bottom):
            return True
        return False

    def find_nearest_car_to_mouse_pos(self, x, y):
        min_distance_car = None

        if self.direction == TO_LEFT or self.direction == TO_RIGHT:
            min_distance = SCREEN_WIDTH
            for car in self.cars:
                if abs(x - car.rect.centerx) < min_distance:
                    min_distance = abs(x - car.rect.centerx)
                    min_distance_car = car

        elif self.direction == TO_BOTTOM or self.direction == TO_TOP:
            min_distance = SCREEN_HEIGHT
            for car in self.cars:
                if abs(y - car.rect.centery) < min_distance:
                    min_distance = abs(y - car.rect.centery)
                    min_distance_car = car

        return min_distance_car

    def include_car(self, car):
        maxtop = max(self.rect.top, car.rect.top)
        minbottom = min(self.rect.bottom, car.rect.bottom)
        if (maxtop < minbottom):
            maxleft = max(self.rect.left, car.rect.left)
            minright = min(self.rect.right, car.rect.right)
            if (maxleft < minright):
                return True
        return False

# Configure lanes before and after, which will be used for changing lanes
def configure_lanes_before_after(lanes):
    for idx, lane in enumerate(lanes):
        if (idx-1 >= 0) and (lanes[idx-1].direction == lane.direction):
            lane.before = lanes[idx-1]
            #print(idx, "before exists")
        if (idx+1 < len(lanes)) and (lanes[idx+1].direction == lane.direction):
            lane.after = lanes[idx+1]
            #print(idx, "after exists")

'''
Define a Car object by extending pygame.sprite.Sprite
The surface drawn on the screen is an attribute of a Car object
'''
class Car():
    def __init__(self, pygame, road, lane, x, y):
        self.road = road
        self.lane = lane

        if self.lane.direction == TO_LEFT or self.lane.direction == TO_RIGHT:
            self.surf = pygame.Surface((CAR_LENGTH, CAR_WIDTH))  # X/Y size
        elif self.lane.direction == TO_BOTTOM or self.lane.direction == TO_TOP:
            self.surf = pygame.Surface((CAR_WIDTH, CAR_LENGTH))  # X/Y size

        self.color = (CAR_COLOR[0]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR),
                        CAR_COLOR[1]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR),
                        CAR_COLOR[2]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR))
        self.surf.fill(self.color)
        self.rect = self.surf.get_rect(center=(x,y))
        self.speed = random.randint(CAR_SPEED-CAR_SPEED_VAR, CAR_SPEED+CAR_SPEED_VAR)
        self.accident = False

    def outside_safe_distance_from(self, car, safe_distance, direction):
        if direction == TO_LEFT or direction == TO_RIGHT:
            if car.rect.right + safe_distance <= self.rect.left or\
               self.rect.right < car.rect.left - safe_distance:
                return True
            else:
                return False
        elif direction == TO_BOTTOM or direction == TO_TOP:
            if car.rect.bottom + safe_distance <= self.rect.top or\
               self.rect.bottom < car.rect.top - safe_distance:
                return True
            else:
                return False

    def paint_on(self, screen):
        screen.blit(self.surf, self.rect)

    '''
    Find the farthest distance that this can can go at the current round,
                        considering other cars and traffic lights
    '''
    def find_farthest_to_go(self):
        if self.lane.direction == TO_LEFT:
            limit = self.lane.x_preceding_car + CAR_SAFE_DISTANCE
            reason = BLOCKED

            # Consider cars on lanes the cross this lane
            for blocking_lane in self.lane.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.lane.include_car(c) and c.rect.right < self.rect.left:
                        limit = max(limit, c.rect.right + CAR_SAFE_DISTANCE)

            # Consider cars and traffic lights on lanes the this lane continues on
            if self.rect.left - self.speed <= self.lane.rect.left + CAR_SAFE_DISTANCE:
                for next_lane in self.lane.next:
                    if next_lane.trafficLight == REDLIGHT:
                        limit = max(limit, next_lane.rect.right + CAR_SAFE_DISTANCE)
                        if limit == next_lane.rect.right + CAR_SAFE_DISTANCE:
                            reason = REDLIGHT
                        break   # At a red light, no need to check other conditions that provide looser limits

                    if len(next_lane.cars) > 0:
                        limit = max(limit, next_lane.cars[-1].rect.right + CAR_SAFE_DISTANCE)

                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c):
                                limit = max(limit, c.rect.right + CAR_SAFE_DISTANCE)

        elif self.lane.direction == TO_RIGHT:
            limit = self.lane.x_preceding_car - CAR_SAFE_DISTANCE
            reason = BLOCKED

            # Consider cars on lanes the cross this lane
            for blocking_lane in self.lane.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.lane.include_car(c) and self.rect.right < c.rect.left:
                        limit = min(limit, c.rect.left - CAR_SAFE_DISTANCE)

            # Consider cars and traffic lights on lanes the this lane continues on
            if self.lane.rect.right - CAR_SAFE_DISTANCE <= self.rect.right + self.speed:
                for next_lane in self.lane.next:
                    if next_lane.trafficLight == REDLIGHT:
                        limit = min(limit, next_lane.rect.left - CAR_SAFE_DISTANCE)
                        if limit == next_lane.rect.left - CAR_SAFE_DISTANCE:
                            reason = REDLIGHT
                        break   # At a red light, no need to check other conditions that provide looser limits

                    if len(next_lane.cars) > 0:
                        limit = min(limit, next_lane.cars[-1].rect.left - CAR_SAFE_DISTANCE)

                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c):
                                limit = min(limit, c.rect.left - CAR_SAFE_DISTANCE)

        elif self.lane.direction == TO_BOTTOM:
            limit = self.lane.y_preceding_car - CAR_SAFE_DISTANCE
            reason = BLOCKED

            # Consider cars on lanes the cross this lane
            for blocking_lane in self.lane.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.lane.include_car(c) and self.rect.bottom < c.rect.top:
                        limit = min(limit, c.rect.top - CAR_SAFE_DISTANCE)

            # Consider cars and traffic lights on lanes the this lane continues on
            if self.lane.rect.bottom - CAR_SAFE_DISTANCE <= self.rect.bottom + self.speed:
                for next_lane in self.lane.next:
                    if next_lane.trafficLight == REDLIGHT:
                        limit = min(limit, next_lane.rect.top - CAR_SAFE_DISTANCE)
                        if limit == next_lane.rect.top - CAR_SAFE_DISTANCE:
                            reason = REDLIGHT
                        break   # At a red light, no need to check other conditions that provide looser limits

                    if len(next_lane.cars) > 0:
                        limit = min(limit, next_lane.cars[-1].rect.top - CAR_SAFE_DISTANCE)

                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c):
                                limit = min(limit, c.rect.top - CAR_SAFE_DISTANCE)

        elif self.lane.direction == TO_TOP:
            limit = self.lane.y_preceding_car + CAR_SAFE_DISTANCE
            reason = BLOCKED

            # Consider cars on lanes the cross this lane
            for blocking_lane in self.lane.blocking_lanes:
                for c in blocking_lane.cars:
                    if self.lane.include_car(c) and c.rect.bottom < self.rect.top:
                        limit = max(limit, c.rect.bottom + CAR_SAFE_DISTANCE)

            # Consider cars and traffic lights on lanes the this lane continues on
            if self.rect.top - self.speed <= self.lane.rect.top + CAR_SAFE_DISTANCE:
                for next_lane in self.lane.next:
                    if next_lane.trafficLight == REDLIGHT:
                        limit = max(limit, next_lane.rect.bottom + CAR_SAFE_DISTANCE)
                        if limit == next_lane.rect.bottom + CAR_SAFE_DISTANCE:
                            reason = REDLIGHT
                        break   # At a red light, no need to check other conditions that provide looser limits

                    if len(next_lane.cars) > 0:
                        limit = max(limit, next_lane.cars[-1].rect.bottom + CAR_SAFE_DISTANCE)

                    for blocking_lane in next_lane.blocking_lanes:
                        for c in blocking_lane.cars:
                            if next_lane.include_car(c):
                                limit = max(limit, c.rect.bottom + CAR_SAFE_DISTANCE)

        return limit, reason

    '''
    Move this car.
    Return True if this car remains on the current lane. Return False otherwise.
    '''
    def move(self, batch):
        if self.lane.direction == TO_LEFT:
            if self.accident:
                self.lane.x_preceding_car = self.rect.right
                self.lane.status_preceding_car = BLOCKED
                return True

            limit, reason = self.find_farthest_to_go()

            prev_lane = self.lane
            prev_left = self.rect.left
            if limit <= self.rect.left - self.speed:    # GO
                self.rect.left = self.rect.left - self.speed    # Move at the assigned speed
                self.lane.x_preceding_car = self.rect.right
                self.lane.status_preceding_car = GO
            else:   # REDLIGHT or BLOCKED
                '''
                At a red light, cars from different lanes may not stop along the same line,
                        because they have come at diffrent speeds, and thus
                        they can see the red light at different positions
                '''
                if (limit < self.rect.left):    # Do not move beyond the limit
                    self.rect.left = limit

                if reason == REDLIGHT:  # REDLIGHT
                    self.lane.x_preceding_car = self.rect.right
                    self.lane.status_preceding_car = REDLIGHT

                else: # BLOCKED
                    if self.change_lane_v2(CAR_CHANGE_LANE_RATE_BLOCKED):
                        pass
                    else:
                        self.lane.x_preceding_car = self.rect.right
                        if self.lane.status_preceding_car == REDLIGHT:
                            pass # Blocked because preceding cars stop at a red light
                        else:
                            # Blocked because of accidents or congestion
                            self.lane.status_preceding_car == BLOCKED
                            if self.rect.left == prev_left: # Could not move at all, thus send a report
                                batch.report(self, report.EVENT_STOP)
                                self.change_color()

            if self.rect.left < self.lane.rect.left: # Add to the next lane
                self.add_to_next_lane()

            if self.lane != prev_lane or self.rect.left < self.lane.rect.left:
                return False    # Lane changed
            else:
                return True     # Lane remains the same

        elif self.lane.direction == TO_RIGHT:
            if self.accident:
                self.lane.x_preceding_car = self.rect.left
                self.lane.status_preceding_car = BLOCKED
                return True

            limit, reason = self.find_farthest_to_go()

            prev_lane = self.lane
            prev_right = self.rect.right
            if self.rect.right + self.speed <= limit:    # GO
                self.rect.right = self.rect.right + self.speed    # Move at the assigned speed
                self.lane.x_preceding_car = self.rect.left
                self.lane.status_preceding_car = GO
            else:   # REDLIGHT or BLOCKED
                '''
                At a red light, cars from different lanes may not stop along the same line,
                        because they have come at diffrent speeds, and thus
                        they can see the red light at different positions
                '''
                if (self.rect.right < limit):    # Do not move beyond the limit
                    self.rect.right = limit

                if reason == REDLIGHT:  # REDLIGHT
                    self.lane.x_preceding_car = self.rect.left
                    self.lane.status_preceding_car = REDLIGHT

                else: # BLOCKED
                    if self.change_lane_v2(CAR_CHANGE_LANE_RATE_BLOCKED):
                        pass
                    else:
                        self.lane.x_preceding_car = self.rect.left
                        if self.lane.status_preceding_car == REDLIGHT:
                            pass # Blocked because preceding cars stop at a red light
                        else:
                            # Blocked because of accidents or congestion
                            self.lane.status_preceding_car == BLOCKED
                            if self.rect.right == prev_right: # Could not move at all, thus send a report
                                batch.report(self, report.EVENT_STOP)
                                self.change_color()

            if self.lane.rect.right < self.rect.right: # Add to the next lane
                self.add_to_next_lane()

            if self.lane != prev_lane or self.lane.rect.right < self.rect.right:
                return False    # Lane changed
            else:
                return True     # Lane remains the same

        elif self.lane.direction == TO_BOTTOM:
            if self.accident:
                self.lane.y_preceding_car = self.rect.top
                self.lane.status_preceding_car = BLOCKED
                return True

            limit, reason = self.find_farthest_to_go()

            prev_lane = self.lane
            prev_bottom = self.rect.bottom
            if self.rect.bottom + self.speed <= limit:    # GO
                self.rect.bottom = self.rect.bottom + self.speed    # Move at the assigned speed
                self.lane.y_preceding_car = self.rect.top
                self.lane.status_preceding_car = GO
            else:   # REDLIGHT or BLOCKED
                '''
                At a red light, cars from different lanes may not stop along the same line,
                        because they have come at diffrent speeds, and thus
                        they can see the red light at different positions
                '''
                if (self.rect.bottom < limit):    # Do not move beyond the limit
                    self.rect.bottom = limit

                if reason == REDLIGHT:  # REDLIGHT
                    self.lane.y_preceding_car = self.rect.top
                    self.lane.status_preceding_car = REDLIGHT

                else: # BLOCKED
                    if self.change_lane_v2(CAR_CHANGE_LANE_RATE_BLOCKED):
                        pass
                    else:
                        self.lane.y_preceding_car = self.rect.top
                        if self.lane.status_preceding_car == REDLIGHT:
                            pass # Blocked because preceding cars stop at a red light
                        else:
                            # Blocked because of accidents or congestion
                            self.lane.status_preceding_car == BLOCKED
                            if self.rect.bottom == prev_bottom: # Could not move at all, thus send a report
                                batch.report(self, report.EVENT_STOP)
                                self.change_color()

            if self.lane.rect.bottom < self.rect.bottom: # Add to the next lane
                self.add_to_next_lane()

            if self.lane != prev_lane or self.lane.rect.bottom < self.rect.bottom:
                return False    # Lane changed
            else:
                return True     # Lane remains the same

        elif self.lane.direction == TO_TOP:
            if self.accident:
                self.lane.y_preceding_car = self.rect.bottom
                self.lane.status_preceding_car = BLOCKED
                return True

            limit, reason = self.find_farthest_to_go()

            prev_lane = self.lane
            prev_top = self.rect.top
            if limit <= self.rect.top - self.speed:    # GO
                self.rect.top = self.rect.top - self.speed    # Move at the assigned speed
                self.lane.y_preceding_car = self.rect.bottom
                self.lane.status_preceding_car = GO
            else:   # REDLIGHT or BLOCKED
                '''
                At a red light, cars from different lanes may not stop along the same line,
                        because they have come at diffrent speeds, and thus
                        they can see the red light at different positions
                '''
                if (limit < self.rect.top):    # Do not move beyond the limit
                    self.rect.top = limit

                if reason == REDLIGHT:  # REDLIGHT
                    self.lane.y_preceding_car = self.rect.bottom
                    self.lane.status_preceding_car = REDLIGHT

                else: # BLOCKED
                    if self.change_lane_v2(CAR_CHANGE_LANE_RATE_BLOCKED):
                        pass
                    else:
                        self.lane.y_preceding_car = self.rect.bottom
                        if self.lane.status_preceding_car == REDLIGHT:
                            pass # Blocked because preceding cars stop at a red light
                        else:
                            # Blocked because of accidents or congestion
                            self.lane.status_preceding_car == BLOCKED
                            if self.rect.top == prev_top: # Could not move at all, thus send a report
                                batch.report(self, report.EVENT_STOP)
                                self.change_color()

            if self.rect.top < self.lane.rect.top: # Add to the next lane
                self.add_to_next_lane()

            if self.lane != prev_lane or self.rect.top < self.lane.rect.top:
                return False    # Lane changed
            else:
                return True     # Lane remains the same

    '''
    Move this car to the lane that the current lane continues on
    '''
    def add_to_next_lane(self):
        if len(self.lane.next) > 0:                     # If a next lane exists, the car continues on the next lane
            # Select a next lane to continue
            if len(self.lane.next) == 1:
                lane = self.lane.next[0]
            else:
                lane = self.lane.next[random.randrange(0,len(self.lane.next)-1)]

            self.road = lane.road
            self.lane = lane
            lane.cars.append(self)
        else:
            pass    # Move outside the screen

    '''
    Switch into lanes before and after the current lane
    '''
    def change_lane_v2(self, probability):
        # Find empty lanes
        lanes = []
        if self.lane.before != None:
            yes, idx = self.lane.before.can_change_lane(self, add=False)
            if yes:
                lanes.append((self.lane.before, idx))
        if self.lane.after != None:
            yes, idx = self.lane.after.can_change_lane(self, add=False)
            if yes:
                lanes.append((self.lane.after, idx))

        # Change lanes probabilistically
        if len(lanes)>0 and random.randrange(1,100) <= (probability * 100):
            if (len(lanes) == 1):
                lane, idx = lanes[0]
            else:
                lane, idx = lanes[random.randrange(0,len(lanes)-1)]
            lane.cars.insert(idx, self)
            self.lane = lane
            return True

        return False

    def change_color(self):
            self.color = (CAR_COLOR[0]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR),
                        CAR_COLOR[1]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR),
                        CAR_COLOR[2]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR))
            self.surf.fill(self.color)

    def toggle_accident(self, batch):
        if not self.accident:
            self.accident = True
            self.prev_color = self.color
            self.color = (CAR_COLOR_ACCIDENT[0]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR),
                        CAR_COLOR_ACCIDENT[1]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR),
                        CAR_COLOR_ACCIDENT[2]+random.randrange(-CAR_COLOR_VAR,CAR_COLOR_VAR))
            self.surf.fill(self.color)
            batch.report(self, report.EVENT_ACCIDENT)

        else:
            self.accident = False
            self.color = self.prev_color
            self.surf.fill(self.color)

    def distance_from(self, x, y):
        return math.sqrt((self.rect.centerx - x)**2 + (self.rect.centery-y)**2)

def find_car_nearest_to_mouse_pos(roads, x, y):
    for road in roads:
        lanes_on_mouse_pos = road.find_lanes_on_mouse_pos(x, y)
        if len(lanes_on_mouse_pos) > 0:
            break

    if len(lanes_on_mouse_pos) > 0:
        cars_on_mouse_pos = []
        for lane in lanes_on_mouse_pos:
            car = lane.find_nearest_car_to_mouse_pos(x, y)
            if car != None:
                cars_on_mouse_pos.append(car)

        if (len(cars_on_mouse_pos) > 0):
            car_nearest_to_mouse_pos = None
            min_distance = math.inf
            for car in cars_on_mouse_pos:
                distance = car.distance_from(x,y)
                if (distance < min_distance):
                    min_distance = distance
                    car_nearest_to_mouse_pos = car
            return car
        else:
            return None

    else:
        return None

'''
Define an Intersection object
'''
class Intersection():
    def __init__(self, left, right, top, bottom, roads):
        self.left = left
        self.right = right
        self.top = top
        self.bottom = bottom
        self.roads = roads

        '''
        surrounding_lanes[0]: lanes on the left side of intersection heading to the left
        surrounding_lanes[1]: lanes on the left side of intersection heading to the right
        surrounding_lanes[2]: lanes on the bottom side of intersection heading to the bottom
        surrounding_lanes[3]: lanes on the bottom side of intersection heading to the top
        surrounding_lanes[4]: lanes on the right side of intersection heading to the left
        surrounding_lanes[5]: lanes on the right side of intersection heading to the right
        surrounding_lanes[6]: lanes on the top side of intersection heading to the bottom
        surrounding_lanes[7]: lanes on the top side of intersection heading to the top
        '''
        self.surrounding_lanes = [ [] for _ in range(8) ]  # Lanes surrounding the intersection
        self.lanes = []                                   # Lanes on the intersection

        self.signal_group = []                     # Group of lanes in the same signal groups
        self.current_signal = 0

        # Added for consistency with Road objects
        self.lanes_for_new_cars = []
        self.name_str = '-'.join([road.name_str for road in self.roads])

    def add_newCar(self):
        pass

    def paint_on(self, screen):
        for lane in self.lanes:
            lane.paint_on(screen)

    def paint_cars_on(self, screen):
        for lane in self.lanes:
            lane.paint_cars_on(screen)

    def move(self, batch):
        for lane in self.lanes:
            lane.move(batch)

    def find_lanes_on_mouse_pos(self, x, y):
        lanes = []
        for lane in self.lanes:
            if lane.include_pos(x, y):
                lanes.append(lane)
        return lanes

    def on_the_same_road_with(self, road):
        if isinstance(road, Road):
            for r in self.roads:
                if r == road:
                    return True
            return False
        elif isinstance(road, Intersection):
            for r1 in self.roads:
                for r2 in road.roads:
                    '''
                    Two intersections are considered being on the same road
                                        if they contain one road in common
                                        because accidents in one intersection can influence the other
                    '''
                    if r1 == r2:
                        return True
            return False
        else:
            return False

    # connect surrounding lanes of the same direction
    def connect_surrounding_lanes_same_direction(self, idx_from, idx_to):
        group = []
        for lanePair in zip(self.surrounding_lanes[idx_from], self.surrounding_lanes[idx_to]):
            from_lane, to_lane = lanePair[0], lanePair[1]
            new_lane = Lane(from_lane.pygame, from_lane.direction, from_lane.position, self)
            if from_lane.direction == TO_LEFT:
                new_lane.update_size(to_lane.rect.right, to_lane.rect.top, self.right - self.left, new_lane.rect.height)
            elif from_lane.direction == TO_RIGHT:
                new_lane.update_size(from_lane.rect.right, from_lane.rect.top, self.right - self.left, new_lane.rect.height)
            elif from_lane.direction == TO_BOTTOM:
                new_lane.update_size(from_lane.rect.left, from_lane.rect.bottom, new_lane.rect.width, self.bottom - self.top)
            elif from_lane.direction == TO_TOP:
                new_lane.update_size(to_lane.rect.left, to_lane.rect.bottom, new_lane.rect.width, self.bottom - self.top)
            new_lane.next.append(to_lane)
            from_lane.next.append(new_lane)
            self.lanes.append(new_lane)
            group.append(new_lane)
        return group

'''
For each pair of intersecting roads R1 and R2,
                    add lanes for switching between roads
'''
def add_intersections(roads):
    '''
    Find all intersections
    '''
    intersections = []
    road_combinations = itertools.combinations(roads, 2)
    for pair in road_combinations:
        pairOrdered = pair[0].intersect(pair[1])
        if pairOrdered != None:
            roadH, roadV = pairOrdered[0], pairOrdered[1]
            left, right, top, bottom = roadV.lanes[0].rect.left, roadV.lanes[-1].rect.right,\
                              roadH.lanes[0].rect.top, roadH.lanes[-1].rect.bottom
            it = Intersection(left, right, top, bottom, [roadH, roadV])
            roadH.intersections.append((left,right,it))
            roadV.intersections.append((top,bottom,it))
            intersections.append(it)
            #print(it.top, it.bottom, it.left, it.right)

    '''
    At each intersection, divide each lane into two,
                        one before the intersection and the other after the intersection
    '''
    for road in roads:
        road.intersections = sorted(road.intersections, key=lambda x: x[0]) # Sort by left or top coordinates to divide from left to right (from top to bottom)
        prev_lanes = road.lanes
        for it in road.intersections:
            #print(road.name_str, it)
            new_lanes = []
            for lane in prev_lanes:
                if road.orientation == HORIZONTAL:
                    lane.update_size(lane.rect.left, lane.rect.top, it[0] - lane.rect.left, lane.rect.height)
                    new_lane = Lane(lane.pygame, lane.direction, lane.position, lane.road)
                    new_lane.update_size(it[1], new_lane.rect.top, new_lane.rect.right - it[1], new_lane.rect.height)
                    if lane.direction == TO_LEFT:
                        it[2].surrounding_lanes[0].append(lane)
                        it[2].surrounding_lanes[5].append(new_lane)
                    elif lane.direction == TO_RIGHT:
                        it[2].surrounding_lanes[1].append(lane)
                        it[2].surrounding_lanes[4].append(new_lane)
                elif road.orientation == VERTICAL:
                    lane.update_size(lane.rect.left, lane.rect.top, lane.rect.width, it[0] - lane.rect.top)
                    new_lane = Lane(lane.pygame, lane.direction, lane.position, lane.road)
                    new_lane.update_size(new_lane.rect.left, it[1], new_lane.rect.width, new_lane.rect.bottom-it[1])
                    if lane.direction == TO_BOTTOM:
                        it[2].surrounding_lanes[7].append(lane)
                        it[2].surrounding_lanes[2].append(new_lane)
                    elif lane.direction == TO_TOP:
                        it[2].surrounding_lanes[6].append(lane)
                        it[2].surrounding_lanes[3].append(new_lane)
                new_lanes.append(new_lane)
                #print(lane.rect.top, lane.rect.bottom, lane.rect.left, lane.rect.right)
                #print(new_lane.rect.top, new_lane.rect.bottom, new_lane.rect.left, new_lane.rect.right)

            configure_lanes_before_after(new_lanes)
            prev_lanes = new_lanes
            road.lanes.extend(new_lanes)
            #print(len(new_lanes),'new lanes added, so a total of', len(road.lanes))

    '''
    At each intersection, connect surrounding lanes
                        with new straight or curved lines
    '''
    for it in intersections:
        # Signal group 1
        group1 = it.connect_surrounding_lanes_same_direction(5,0)
        add_blocking_lanes(group1, it.surrounding_lanes[6])
        group1a = it.connect_surrounding_lanes_same_direction(1,4)
        add_blocking_lanes(group1a, it.surrounding_lanes[2])

        group1.extend(group1a)
        configure_lanes_before_after(group1)
        it.signal_group.append(group1)

        # Signal group 2
        group2 = it.connect_surrounding_lanes_same_direction(7,2)
        add_blocking_lanes(group2, it.surrounding_lanes[0])
        group2a = it.connect_surrounding_lanes_same_direction(3,6)
        add_blocking_lanes(group2a, it.surrounding_lanes[4])

        group2.extend(group2a)
        configure_lanes_before_after(group2)
        it.signal_group.append(group2)

        # Configure blocking lanes
        add_blocking_lanes(group1, group2)
        add_blocking_lanes(group2, group1)

    '''
    At each intersection, randomly order signal groups
    '''
    for it in intersections:
        random.shuffle(it.signal_group)
        for idx, group in enumerate(it.signal_group):
            if idx == it.current_signal:
                for lane in group:
                    lane.trafficLight = GO
            else:
                for lane in group:
                    lane.trafficLight = REDLIGHT

    roads.extend(intersections)
    '''
    for road in roads:
        print(road.name_str)
        for lane in road.lanes:
            print(lane.rect.top, lane.rect.bottom, lane.rect.left, lane.rect.right)
    '''
    return intersections

def add_blocking_lanes(to_lanes, blocking_lanes):
    for lane in to_lanes:
        lane.blocking_lanes.extend(blocking_lanes)
        #print(lane.rect, len(lane.blocking_lanes))


import logging
from . import epdconfig

# Display resolution
EPD_WIDTH       = 640
EPD_HEIGHT      = 400

logger = logging.getLogger(__name__)

class EPD:
    def __init__(self):
        self.reset_pin = epdconfig.RST_PIN
        self.dc_pin = epdconfig.DC_PIN
        self.busy_pin = epdconfig.BUSY_PIN
        self.cs_pin = epdconfig.CS_PIN
        self.width = EPD_WIDTH
        self.height = EPD_HEIGHT
        self.BLACK  = 0x000000   #   0000  BGR
        self.WHITE  = 0xffffff   #   0001
        self.GREEN  = 0x00ff00   #   0010
        self.BLUE   = 0xff0000   #   0011
        self.RED    = 0x0000ff   #   0100
        self.YELLOW = 0x00ffff   #   0101
        self.ORANGE = 0x0080ff   #   0110

    # Hardware reset
    def reset(self):
        epdconfig.digital_write(self.reset_pin, 1)
        epdconfig.delay_ms(200)
        epdconfig.digital_write(self.reset_pin, 0)
        epdconfig.delay_ms(1)
        epdconfig.digital_write(self.reset_pin, 1)
        epdconfig.delay_ms(200)

    def send_command(self, command):
        epdconfig.digital_write(self.dc_pin, 0)
        epdconfig.digital_write(self.cs_pin, 0)
        epdconfig.spi_writebyte([command])
        epdconfig.digital_write(self.cs_pin, 1)

    def send_data(self, data):
        epdconfig.digital_write(self.dc_pin, 1)
        epdconfig.digital_write(self.cs_pin, 0)
        epdconfig.spi_writebyte([data])
        epdconfig.digital_write(self.cs_pin, 1)

    # send a lot of data
    def send_data2(self, data):
        epdconfig.digital_write(self.dc_pin, 1)
        epdconfig.digital_write(self.cs_pin, 0)
        epdconfig.spi_writebyte2(data)
        epdconfig.digital_write(self.cs_pin, 1)

    def ReadBusyHigh(self):
        logger.debug("e-Paper busy")
        while(epdconfig.digital_read(self.busy_pin) == 0):      # 0: idle, 1: busy
            epdconfig.delay_ms(10)
        logger.debug("e-Paper busy release")

    def ReadBusyLow(self):
        logger.debug("e-Paper busy")
        while(epdconfig.digital_read(self.busy_pin) == 1):      # 0: idle, 1: busy
            epdconfig.delay_ms(10)
        logger.debug("e-Paper busy release")

    def init(self):
        if (epdconfig.module_init() != 0):
            return -1
        # EPD hardware init start
        self.reset()

        self.ReadBusyHigh()
        self.send_command(0x00)
        self.send_data(0x2f)
        self.send_data(0x00)
        self.send_command(0x01)
        self.send_data(0x37)
        self.send_data(0x00)
        self.send_data(0x05)
        self.send_data(0x05)
        self.send_command(0x03)
        self.send_data(0x00)
        self.send_command(0x06)
        self.send_data(0xC7)
        self.send_data(0xC7)
        self.send_data(0x1D)
        self.send_command(0x41)
        self.send_data(0x00)
        self.send_command(0x50)
        self.send_data(0x37)
        self.send_command(0x60)
        self.send_data(0x22)
        self.send_command(0x61)
        self.send_data(0x02)
        self.send_data(0x80)
        self.send_data(0x01)
        self.send_data(0x90)
        self.send_command(0xE3)
        self.send_data(0xAA)

        # EPD hardware init end
        return 0

    def getbuffer(self, image):
        buf = [0x00] * int(self.width * self.height / 2)
        image_monocolor = image.convert('RGB')#Picture mode conversion
        imwidth, imheight = image_monocolor.size
        pixels = image_monocolor.load()
        logger.debug('imwidth = %d  imheight =  %d ',imwidth, imheight)
        if(imwidth == self.width and imheight == self.height):
            for y in range(imheight):
                for x in range(imwidth):
                    # Set the bits for the column of pixels at the current position.
                    Add = int((x + y * self.width) / 2)
                    Color = 0;
                    if (pixels[x, y][0] == 0 and  pixels[x, y][1] == 0 and pixels[x, y][2] == 0):
                        Color = 0
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 255 and pixels[x, y][2] == 255):
                        Color = 1
                    elif (pixels[x, y][0] == 0 and  pixels[x, y][1] == 255 and pixels[x, y][2] == 0):
                        Color = 2
                    elif (pixels[x, y][0] == 0 and  pixels[x, y][1] == 0 and pixels[x, y][2] == 255):
                        Color = 3
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 0 and pixels[x, y][2] == 0):
                        Color = 4
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 255 and pixels[x, y][2] == 0):
                        Color = 5
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 128 and pixels[x, y][2] == 0):
                        Color = 6

                    data_t = buf[Add]&(~(0xF0 >> ((x % 2)*4)))
                    buf[Add] = data_t | ((Color << 4) >> ((x % 2)*4));

        elif(imwidth == self.height and imheight == self.width):
            for y in range(imheight):
                for x in range(imwidth):
                    newx = y
                    newy = self.height - x - 1
                    Add = int((newx + newy*self.width) / 2)
                    Color = 0;
                    if (pixels[x, y][0] == 0 and  pixels[x, y][1] == 0 and pixels[x, y][2] == 0):
                        Color = 0
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 255 and pixels[x, y][2] == 255):
                        Color = 1
                    elif (pixels[x, y][0] == 0 and  pixels[x, y][1] == 255 and pixels[x, y][2] == 0):
                        Color = 2
                    elif (pixels[x, y][0] == 0 and  pixels[x, y][1] == 0 and pixels[x, y][2] == 255):
                        Color = 3
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 0 and pixels[x, y][2] == 0):
                        Color = 4
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 255 and pixels[x, y][2] == 0):
                        Color = 5
                    elif (pixels[x, y][0] == 255 and  pixels[x, y][1] == 128 and pixels[x, y][2] == 0):
                        Color = 6

                    data_t = buf[Add]&(~(0xF0 >> ((newx % 2)*4)))
                    buf[Add] = data_t | ((Color << 4) >> ((newx % 2)*4));
        return buf

    def display(self,image):
        self.send_command(0x61)#Set Resolution setting
        self.send_data(0x02)
        self.send_data(0x80)
        self.send_data(0x01)
        self.send_data(0x90)
        self.send_command(0x10)
        self.send_data2(image)
        self.send_command(0x04)#0x04
        self.ReadBusyHigh()
        self.send_command(0x12)#0x12
        self.ReadBusyHigh()
        self.send_command(0x02)  #0x02
        self.ReadBusyLow()
        # epdconfig.delay_ms(500)

    def Clear(self):
        self.send_command(0x61)#Set Resolution setting
        self.send_data(0x02)
        self.send_data(0x80)
        self.send_data(0x01)
        self.send_data(0x90)
        self.send_command(0x10)
        self.send_data2([0x11] * int(EPD_HEIGHT) * int(EPD_WIDTH/2))
        #BLACK   0x00    /// 0000
        #WHITE   0x11    /// 0001
        #GREEN   0x22    /// 0010
        #BLUE    0x33    /// 0011
        #RED     0x44    /// 0100
        #YELLOW  0x55    /// 0101
        #ORANGE  0x66    /// 0110
        #CLEAN   0x77    /// 0111   unavailable  Afterimage
        self.send_command(0x04)#0x04
        self.ReadBusyHigh()
        self.send_command(0x12)#0x12
        self.ReadBusyHigh()
        self.send_command(0x02)  #0x02
        self.ReadBusyLow()
        # epdconfig.delay_ms(500)

    def sleep(self):
        # epdconfig.delay_ms(500)
        self.send_command(0x07) # DEEP_SLEEP
        self.send_data(0XA5)

        epdconfig.delay_ms(2000)
        epdconfig.module_exit()


import shelve
import os
import csv
import urllib
from functools import reduce

from luigi import format
from luigi import Parameter
from luigi import Task
from luigi import WrapperTask
from luigi.task_register import load_task
from luigi.task import logger as luigi_logger

import requests
from lxml import html
import pandas as pd

from ..utils import bytes_pls
from ..utils import clean
from ..utils import sha256sum
from ..utils import TargetOutput
from ..utils import Requires
from ..utils import Requirement
from ..utils import SuffixPreservingLocalTarget as LocalTarget

USE_SHELVE = False
DEV_MODE = False

REMOTE_FILE_SOURCES = {
    'unterm': 'https://protocol.un.org/dgacm/pls/site.nsf/files/Country%20Names%20UNTERM2/$FILE/UNTERM%20EFSRCA.xlsx',
    'iso4217': 'https://www.currency-iso.org/dam/downloads/lists/list_one.xml',
    'marc': 'http://www.loc.gov/standards/codelists/countries.xml',
    'cldr': 'https://raw.githubusercontent.com/unicode-cldr/cldr-localenames-full/master/main/en/territories.json',
    'geonames': 'http://download.geonames.org/export/dump/countryInfo.txt',
    'usa-census': 'https://www.census.gov/foreign-trade/schedules/c/country2.txt',
    'exio-wiod-eora': 'https://raw.githubusercontent.com/konstantinstadler/country_converter/master/country_converter/country_data.tsv',
    'ukgov': 'https://country.register.gov.uk/records.json'
}

CUSTOM_SCRAPE_SOURCES = {
    'Edgar': 'https://www.sec.gov/edgar/searchedgar/edgarstatecodes.htm',
    'M49': 'https://unstats.un.org/unsd/methodology/m49/overview/',
}

SIMPLE_TABLE_SCRAPE_SOURCES = {
    'itu-glad': 'https://www.itu.int/gladapp/GeographicalArea/List',
    'fao': 'http://www.fao.org/countryprofiles/iso3list/en/',
    'fifa-ioc': 'https://simple.wikipedia.org/wiki/Comparison_of_IOC,_FIFA,_and_ISO_3166_country_codes',
}

SOURCES = [
    {
        'name': 'unterm',
        'title': 'United Nations Protocol and Liason Service',
        'path': 'https://protocol.un.org/dgacm/pls/site.nsf/files/Country%20Names%20UNTERM2/$FILE/UNTERM%20EFSRCA.xlsx',
    },
    {
        'name': 'iso4217',
        'title': 'Swiss Association for Standardization',
        'path': 'https://www.currency-iso.org/dam/downloads/lists/list_one.xml',
    },
    {
        'name': 'marc',
        'title': 'USA Library of Congress',
        'path': 'http://www.loc.gov/standards/codelists/countries.xml',
    },
    {
        'name': 'cldr',
        'title': 'Unicode Common Locale Data Repository',
        'path': 'https://raw.githubusercontent.com/unicode-cldr/cldr-localenames-full/master/main/en/territories.json',
    },
    {
        'name': 'geonames',
        'title': 'GeoNames',
        'path': 'http://download.geonames.org/export/dump/countryInfo.txt',
    },
    {
        'name': 'usa-census',
        'title': 'USA Census Bureau',
        'path': 'https://www.census.gov/foreign-trade/schedules/c/country2.txt',
    },
    {
        'name': 'exio-wiod-eora',
        'title': 'Secondary source providing EXIOBASE, WIOD, Eora, and other country codes',
        'path': 'https://raw.githubusercontent.com/konstantinstadler/country_converter/master/country_converter/country_data.tsv',
    },
    {
        'name': 'ukgov',
        'title': 'Government of the United Kingdom',
        'path': 'https://country.register.gov.uk/records.json'
    },
    {
        'name': 'Edgar',
        'title': 'USA Security and Exchange Commission',
        'path': 'https://www.sec.gov/edgar/searchedgar/edgarstatecodes.htm',
    },
    {
        'name': 'M49',
        'title': 'United Nations Statistics Division',
        'path': 'https://unstats.un.org/unsd/methodology/m49/overview/',
    },
    {
        'name': 'itu-glad',
        'title': 'International Telecommunications Union',
        'path': 'https://www.itu.int/gladapp/GeographicalArea/List',
    },
    {
        'name': 'fao',
        'title': 'Food and Agriculture Organization',
        'path': 'http://www.fao.org/countryprofiles/iso3list/en/',
    },
    {
        'name': 'fifa-ioc',
        'title': 'Wikipedia (FIFA and IOC)',
        'path': 'https://simple.wikipedia.org/wiki/Comparison_of_IOC,_FIFA,_and_ISO_3166_country_codes',
    },
]

class FileSource(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'

    slug = Parameter()
    ext = Parameter()

    pattern = '{task.__class__.__name__}-{task.slug}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget,
                          format=format.Nop)

    def run(self):
        url = REMOTE_FILE_SOURCES.get(self.slug)

        with requests.get(url, stream=True) as r:
            r.raise_for_status()
            if r.encoding is None:
                r.encoding = 'utf-8'
            with self.output().open('w') as f:
                for chunk in r.iter_content(chunk_size=128):
                    f.write(bytes_pls(chunk))

class SaltedFileSource(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'

    slug = Parameter()
    ext = Parameter()
    salt = sha256sum()

    requires = Requires()
    source = Requirement(FileSource, slug=slug, ext=ext)

    pattern = '{task.__class__.__name__}-{task.slug}-{task.salt}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget,
                          format=format.Nop)

    def run(self):
        source = self.requires().get('source').output()
        target = self.output().path
        source.copy(target)

class EdgarSource(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'
    slug = Parameter(default='Edgar')
    ext = Parameter(default='.csv')

    pattern = '{task.__class__.__name__}-{task.slug}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget)

    def run(self):
        url = CUSTOM_SCRAPE_SOURCES.get('Edgar')

        content = urllib.request.urlopen(url).read()
        doc = html.fromstring(content)
        rows = doc.xpath('//table')[3].getchildren()

        seen_other_countries = False
        data = []

        for row in rows:
            if seen_other_countries is not True:
                if clean(row.text_content()) != 'Other Countries':
                    continue
                else:
                    seen_other_countries = True
                    continue

            cells = row.getchildren()
            if len(cells) != 2:
                luigi_logger.debug('ERROR IN CELL COUNT')
                for cell in cells:
                    luigi_logger.debug(cell)
                    luigi_logger.debug(cell.text_content())
                continue
            code = clean(cells[0].text_content())
            name = clean(cells[1].text_content())
            data.append((code, name))

        with self.output().open('w') as f:
            writer = csv.writer(f)
            writer.writerows(data)

class SaltedEdgarSource(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'

    slug = Parameter(default='Edgar')
    ext = Parameter(default='.csv')
    salt = sha256sum()

    requires = Requires()
    source = Requirement(EdgarSource, slug='Edgar', ext='.csv')

    pattern = '{task.__class__.__name__}-{task.slug}-{task.salt}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget)

    def run(self):
        self.requires().get('source').output().copy(self.output().path)

class M49Source(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'
    slug = Parameter(default='M49')
    ext = Parameter(default='.csv')

    pattern = '{task.__class__.__name__}-{task.slug}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget)

    def run(self):
        url = CUSTOM_SCRAPE_SOURCES.get('M49')

        if USE_SHELVE:
            luigi_logger.debug('USING SHELVE')
            try:
                shelf = shelve.open('tmpdb')
                content = shelf['M49']
            except KeyError:
                content = urllib.request.urlopen(url).read()
                shelf['M49'] = content
            finally:
                shelf.close()
        else:
            content = urllib.request.urlopen(url).read()

        tables = {'en': 'downloadTableEN', 'cn': 'downloadTableZH',
                  'ru': 'downloadTableRU', 'fr': 'downloadTableFR',
                  'es': 'downloadTableES', 'ar': 'downloadTableAR'}

        def read_table(table_language, table_id):
            l = table_language
            attrs = {'id': table_id}
            converters = {
                    'Region Code': lambda x: str(x),
                    'Intermediate Region Code': lambda x: str(x),
                    'Least Developed Countries (LDC)': lambda x: bool(x) if bool(x) else '',
                    'Land Locked Developing Countries (LLDC)': lambda x: bool(x) if bool(x) else '',
                    'Small Island Developing States (SIDS)': lambda x: bool(x) if bool(x) else '',
                    'Developed /  Developing Countries': lambda x: x.replace('\r\n', ''),
                    }
            df = pd.read_html(content, encoding='utf-8', attrs=attrs,
                              header=0, converters=converters)
            return df[0]

        if USE_SHELVE:
            luigi_logger.debug('USING SHELVE')
            try:
                shelf = shelve.open('tmpdb')
                frames_tuples = shelf['m49-frames_tuples']
            except KeyError:
                # read the 6 html tables into dataframes,
                # arrange dataframes in list of 2-item tuples
                # like [('language', dataframe),...]
                frames_tuples = [(lang, read_table(lang, table)) for lang, table in tables.items()]
                shelf['m49-frames_tuples'] = frames_tuples
            finally:
                shelf.close()
        else:
            frames_tuples = [(lang, read_table(lang, table)) for lang, table in tables.items()]

        # values in these columns are the same in any language
        # (excluding `M49 Code` bc we need it to merge dataframes)
        non_local = ['Global Code', 'Region Code', 'Sub-region Code',
                     'Intermediate Region Code',
                     'ISO-alpha3 Code', 'Least Developed Countries (LDC)',
                     'Land Locked Developing Countries (LLDC)',
                     'Small Island Developing States (SIDS)',
                     'Developed / Developing Countries']

        # remove repeated, non-localized columns
        # from all dataframes except one
        pruned = [(lang, frame.drop([c for c in non_local],
                                     inplace=True, axis=1))
                   for lang, frame in frames_tuples if lang != 'en']

        # reduce frames_tuples:
        # - merge dataframes on 'M49 Code' column
        # - add language suffixes to column names
        merged = reduce(lambda left, right:
                        (right[0], pd.merge(left[1], right[1],
                                            on='M49 Code',
                                            suffixes=('_' + left[0], '_' + right[0]))),
                        frames_tuples)

        with self.output().open('w') as f:
            merged[1].to_csv(f, index=False)

class SaltedM49Source(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'

    slug = Parameter(default='M49')
    ext = Parameter(default='.csv')
    salt = sha256sum()

    requires = Requires()
    source = Requirement(M49Source, slug='M49', ext='.csv')

    pattern = '{task.__class__.__name__}-{task.slug}-{task.salt}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget)

    def run(self):
        self.requires().get('source').output().copy(self.output().path)

class SimpleTableScrapeSource(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'
    slug = Parameter()
    ext = Parameter(default='.csv')

    pattern = '{task.__class__.__name__}-{task.slug}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget)

    def run(self):
        url = SIMPLE_TABLE_SCRAPE_SOURCES.get(self.slug)
        if USE_SHELVE:
            luigi_logger.debug('USING SHELVE')
            try:
                shelf = shelve.open('tmpdb')
                content = shelf[self.slug]
            except KeyError:
                content = urllib.request.urlopen(url).read()
                shelf[self.slug] = str(content)
            finally:
                shelf.close()
        else:
            content = urllib.request.urlopen(url).read()

        df = pd.read_html(content, header=0)
        with self.output().open('w') as f:
            df[0].to_csv(f, index=False)

class SaltedSTSSource(Task):
    __version__ = '0.1'
    DATA_ROOT = 'build/'

    slug = Parameter()
    ext = Parameter(default='.csv')
    salt = sha256sum()

    requires = Requires()
    source = Requirement(SimpleTableScrapeSource, slug=slug, ext=ext)

    pattern = '{task.__class__.__name__}-{task.slug}-{task.salt}{task.ext}'
    output = TargetOutput(file_pattern=pattern, ext='',
                          base_dir=DATA_ROOT,
                          target_class=LocalTarget)

    def run(self):
        self.requires().get('source').output().copy(self.output().path)

class SaltedSources(WrapperTask):

    def requires(self):
        for slug, url in REMOTE_FILE_SOURCES.items():
            _, ext = os.path.splitext(url)
            luigi_logger.debug(['SaltedFileSource', slug, ext])
            yield SaltedFileSource(slug=slug, ext=ext)

        for slug, url in SIMPLE_TABLE_SCRAPE_SOURCES.items():
            luigi_logger.debug(['SaltedSTSSource', slug, '.csv'])
            yield SaltedSTSSource(slug=slug, ext='.csv')

        for slug, url in CUSTOM_SCRAPE_SOURCES.items():
            luigi_logger.debug(['Salted__Source', slug, '.csv'])
            yield load_task(__name__, f'Salted{slug}Source', {'slug': slug, 'ext': '.csv'})


import bpy
from bpy.types import AddonPreferences, PropertyGroup
from .update import *

from bpy.props import (
    IntProperty,
    FloatProperty,
    BoolProperty,
    StringProperty,
    PointerProperty,
    CollectionProperty,
    EnumProperty,
)

def GetLocationPresets(scene, context):

    items = [
        ('0', "None", "", 0),
    ]

    try:
        addon_prefs = context.preferences.addons[__package__].preferences
        file_data = bpy.data.objects[addon_prefs.pak_filedata_name].PAK_FileData

    except KeyError:
        return items

    for i,x in enumerate(file_data.locations):
        items.append((str(i+1), x.name, x.name, i+1))

    return items

def GetImageFormats(scene, context):

    items = [
        ('0', "None", "", 0),
    ]

    try:
        addon_prefs = context.preferences.addons[__package__].preferences
        file_data = bpy.data.objects[addon_prefs.pak_filedata_name].PAK_FileData

    except KeyError:
        return items

    return

class PAK_Image(PropertyGroup):
    """
    Define an object as a list property, for use when displaying objects
    in the user interface.
    """

    enable_export: BoolProperty(
        name = "Enable Export",
        description = "Enable or disable the texture for export",
        default = False,
    )

    export_location: EnumProperty(
        name = "Export Location",
        description = "Set the file path that the texture will be exported to",
        items = GetLocationPresets,
    )

class PAK_ImageItem(PropertyGroup):
    """
    Defines a single texture within a PAK_ImageBundle.
    It just wraps an image type in a PropertyGroup as Blender is a bit dumb with collections.
    """

    tex: PointerProperty(
        type = bpy.types.Image,
        name = "Texture",
        description = "A pointer for the texture this list item represents",
    )

class PAK_ImageBundle(PropertyGroup):
    """
    Define an object as a list property, for use when displaying objects in the user interface
    """

    bundle_items: CollectionProperty(type = PAK_ImageItem)

    name: StringProperty(
        name = "Texture/Bundle Name",
        description = "",
        default = "",
        update = PAK_Update_TextureListItem_Name
    )

    enable_export: BoolProperty(
        name = "Enable Export",
        description = "Enable or disable the texture for export",
        default = False,
        update = PAK_Update_TextureListItem_EnableExport,
    )

    export_location: EnumProperty(
        name = "Export Location",
        description = "Set the file path that the texture will be exported to",
        update = PAK_Update_TextureListItem_ExportLocation,
        items = GetLocationPresets,

    )

    # Used for multi-select mode.
    is_selected: BoolProperty(
        default = False,
        update = PAK_Update_TextureListItem_IsSelected,
    )

class PAK_ExportLocations(PropertyGroup):
    # Defines a single export location that can be reused across multiple textures

    name: StringProperty(
        name = "",
        description = "The name of the export location"
    )

    path: StringProperty(
        name = "",
        description = "The file path used to export images to",
        default = "",
        subtype = "FILE_PATH"
    )

class PAK_FileData(PropertyGroup):
    """
    Everything PakPal needs to preserve as part of the file.
    """

    # if true, this object is the empty created for the purposes of storing preset data.
    is_file_data: BoolProperty(default = False)

    # the available images for export
    bundles: CollectionProperty(type = PAK_ImageBundle)

    ## The index of the currently selected collection from the UI list.  Will be -1 if not selected.
    bundles_list_index: IntProperty(
        default = 0,
        update = PAK_Update_TextureList_Preview
    )

    # the available baking presets
    locations: CollectionProperty(type = PAK_ExportLocations)

    ## The index of the currently selected collection from the UI list.  Will be -1 if not selected.
    locations_list_index: IntProperty(default = 0)

    case_sensitive_matching: BoolProperty(
        name = "Case Sensitive Matching",
        description = "If enabled, operations that use material slot names such as image packing and material bundling will perform material slot comparisons in a case sensitive manner"
    )

    # The scene used to store Compositor nodes in.
    scene_data: PointerProperty(
        type = bpy.types.Scene,
        name = "PakPal Scene Data Source",
        description = "Defines the scene used to store Compositor nodes used for managing Image Format settings"
    )

    ## HIDDEN PROPERTIES
    # (these are activated with operators due to styling issues)

    # if true, images can be edited with a psuedo multiselect interface.
    enable_multiselect: BoolProperty(
        name = "Enable Multiselect",
        description = "When enabled, multiple textures can be selected and edited at the same time",
        default = False,
    )

    # if true, images can be edited with a psuedo multiselect interface.
    enable_bundles: BoolProperty(
        name = "Enable Bundles",
        description = "When enabled, textures sharing the same name but using common suffixes (like BaseColor, Height, etc) will be represented as a single entry in the list, and can be edited as a set of images",
        default = False,
        update = PAK_Update_EnableBundles,
    )

    # if true, show hidden images in the list
    enable_hidden: BoolProperty(
        name = "Show Hidden Images",
        description = "Enable or disable the ability to see hidden image blocks.  These might be used by other addons so be careful!",
        default = False,
    )

    preview_tex: PointerProperty(
        type = bpy.types.Texture,
        name = "Preview Texture",
        description = "Used to preview images in the interface"
    )

    preview_rgb: EnumProperty(
        name = "RGB Preview",
        options = {'ENUM_FLAG'},
        items = (('R', "", "", 'COLOR_RED', 1),
                ('G', "", "", 'COLOR_GREEN', 2),
                ('B', "", "", 'COLOR_BLUE', 100), # due to weird behaviour
        ),
        description = "Change how the preview displays different color channels (alpha display is currently unsupported due to inbuilt image preview limitations)",
        default = {'R', 'G', 'B'},
        update = PAK_Update_TextureList_PreviewColor
    )

    # ////////////////////////////////////////////////////////////////////
    # ////////////////////////////////////////////////////////////////////
    ## PROXIES
    # These are used for multi-select operations and UI enhancements.

    # Used during a refresh to prevent unnecessary update function triggers.
    is_internal_update: BoolProperty(default = False)

    proxy_enable_export: BoolProperty(
        name = "Enable Export",
        description = "Enable or disable the texture for export",
        default = False,
        update = PAK_Update_EnableExport,
    )

    proxy_export_location: EnumProperty(
        name = "Export Location",
        description = "Set the file path that the texture will be exported to",
        items = GetLocationPresets,
        update = PAK_Update_ExportLocation,
    )

    # ////////////////////////////////////////////////////////////////////
    # ////////////////////////////////////////////////////////////////////
    # IMAGE PACKING
    # These options are used when performing a channel mix operation
    pack_r_source: StringProperty(
        name = "Red Source Slot Name",
        description = "Set the name of the material slot that will be used as a source for the new image's red channel (if it can be found within a bundle).  Multiple slot names can be defined but only the first one found in a image bundle will be used",
        default = ""
    )

    pack_g_source: StringProperty(
        name = "Green Source Slot Name",
        description = "Set the name of the material slot that will be used as a source for the new image's green channel (if it can be found within a bundle).  Multiple slot names can be defined but only the first one found in a image bundle will be used",
        default = ""
    )

    pack_b_source: StringProperty(
        name = "Blue Source Slot Name",
        description = "Set the name of the material slot that will be used as a source for the new image's blue channel (if it can be found within a bundle).  Multiple slot names can be defined but only the first one found in a image bundle will be used",
        default = ""
    )

    pack_a_source: StringProperty(
        name = "Alpha Source Slot Name",
        description = "Set the name of the material slot that will be used as a source for the new image's alpha channel (if it can be found within a bundle).  Multiple slot names can be defined but only the first one found in a image bundle will be used",
        default = ""
    )

    pack_r_channel: EnumProperty(
		name = " ",
		items = (('R', "R", "Use the red channel from the source material slot as the new packed image's red channel"),
			    ('G', "G", "Use the green channel from the source material slot as the new packed image's red channel"),
			    ('B', "B", "Use the blue channel from the source material slot as the new packed image's red channel"),
                ('A', "A", "Use the alpha channel from the source material slot as the new packed image's red channel")),
		description = "",
		default = 'R',
	)

    pack_g_channel: EnumProperty(
		name = " ",
		items = (('R', "R", "Use the red channel from the source material slot as the new packed image's green channel"),
			    ('G', "G", "Use the green channel from the source material slot as the new packed image's green channel"),
			    ('B', "B", "Use the blue channel from the source material slot as the new packed image's green channel"),
                ('A', "A", "Use the alpha channel from the source material slot as the new packed image's green channel")),
		description = "",
		default = 'R',
	)

    pack_b_channel: EnumProperty(
		name = " ",
		items = (('R', "R", "Use the red channel from the source material slot as the new packed image's blue channel"),
			    ('G', "G", "Use the green channel from the source material slot as the new packed image's blue channel"),
			    ('B', "B", "Use the blue channel from the source material slot as the new packed image's blue channel"),
                ('A', "A", "Use the alpha channel from the source material slot as the new packed image's blue channel")),
		description = "",
		default = 'R',
	)

    pack_a_channel: EnumProperty(
		name = " ",
		items = (('R', "R", "Use the red channel from the source material slot as the new packed image's alpha channel"),
			    ('G', "G", "Use the green channel from the source material slot as the new packed image's alpha channel"),
			    ('B', "B", "Use the blue channel from the source material slot as the new packed image's alpha channel"),
                ('A', "A", "Use the alpha channel from the source material slot as the new packed image's alpha channel")),
		description = "",
		default = 'R',
	)

    pack_r_invert: BoolProperty(
        name = "Invert R Source",
        description = "Invert the source image output",
        default = False,
    )

    pack_g_invert: BoolProperty(
        name = "Invert G Source",
        description = "Invert the source image output",
        default = False,
    )

    pack_b_invert: BoolProperty(
        name = "Invert B Source",
        description = "Invert the source image output",
        default = False,
    )

    pack_a_invert: BoolProperty(
        name = "Invert A Source",
        description = "Invert the source image output",
        default = False,
    )

    packed_image_suffix: StringProperty(
        name = "Packed Image Suffix",
        description = "Set the suffix the new image will be given.  The base name of the new image will be the same as the base name for the Bundle",
        default = "",
    )
    overwrite_image_pack: BoolProperty(
        name = "Overwrite Existing Image",
        description = "If an image is found with the same name in the Blend file, the new packed image will replace it",
        default = True,
    )

    add_fake_user: BoolProperty(
        name = "Add Fake User",
        description = "Adds a fake user to any generated packed images to prevent them from disappearing when the blend file is closed",
        default = True,
    )

    temp_bake_path: StringProperty(
        name = "Temp Save Location",
        description = "A temporary location used to save packed images.  Either due to Blender limitations or a skill issue, PakPal needs to save images somewhere before reloading and saving them in the blend file",
        default = "//Pak_Cache\\",
        subtype = "FILE_PATH"
    )

class PAK_MaterialSlot(PropertyGroup):

    text: StringProperty(
        name = "Bundle Text",
        description = "",
        default = "",
    )

# As of 3.5 PointerProperty types currently only apply to ID and PropertyGroup subclasses
# If this changes in a future version, replace this!
# MAINTENANCE : This will need to be checked regularly.
# https://docs.blender.org/api/current/bpy.types.ImageFormatSettings.html#bpy.types.ImageFormatSettings
class PAK_ImageFormat(PropertyGroup):

    cineon_black: IntProperty()
    cineon_gamma: FloatProperty()
    cineon_white: IntProperty()

    color_depth: EnumProperty(
        items = (('8', "8", ""),
			    ('10', "10", ""),
			    ('12', "12", ""),
                ('16', "16", ""),
                ('32', "32", "")
                ),
    )

    # can I just use a string property?  OwO
    color_management: StringProperty()
    color_mode: StringProperty()

    compression: IntProperty()

    # (readonly)
    # display_settings - "Settings of devioce saved image would be displayed on"

    # more fake enums
    exr_codec: StringProperty()
    file_format: StringProperty()

    has_linear_colorspace: BoolProperty()

    # fake enum
    jpeg2k_codec: StringProperty()

    # (readonly)
    # linear_colorspace_settings -= "Output color space settings", opaque type.

    quality: IntProperty()
    # (readonly)
    # stereo_3d_format - Another weird structure

    tiff_codec: StringProperty() # fake enum

    use_cineon_log: BoolProperty()

    use_jpeg2k_cinema_48: BoolProperty()
    use_jpeg2k_cinema_preset: BoolProperty()
    use_jpeg2k_ycc: BoolProperty()

    use_preview: BoolProperty() # saves JPG images of animations to the same directory

    # (readonly)
    # view_settings - appears to be the color management settings we're interested in
    # (readonly)
    # views_format - For stereo output, we don't care.


import git
import sys

git_root = git.Repo("", search_parent_directories=True).git.rev_parse("--show-toplevel")
sys.path.append(git_root)

from models.components.modules import *

from tools.pytorch.losses import *
from tools.pytorch.pytorch_util import *
from tools.pytorch.modules.fDomainHelper import FDomainHelper

def to_log(input):
    return torch.log10(torch.clip(input, min=1e-8))

def from_log(input):
    input = torch.clip(input,min=-np.inf, max=5)
    return 10 ** input

class UNetResComplex_100Mb(nn.Module):
    def __init__(self, channels, nsrc=1):
        super(UNetResComplex_100Mb, self).__init__()
        window_size = 2048
        hop_size = 441
        activation = 'relu'
        momentum = 0.01
        center = True,
        pad_mode = 'reflect'
        window = 'hann'
        freeze_parameters = True

        self.nsrc = nsrc
        self.channels = channels
        self.time_downsample_ratio = 2 ** 6  # This number equals 2^{#encoder_blcoks}

        self.f_helper = FDomainHelper(
            window_size=window_size,
            hop_size=hop_size,
            center=center,
            pad_mode=pad_mode,
            window=window,
            freeze_parameters=freeze_parameters,
        )

        self.encoder_block1 = EncoderBlockRes4B(in_channels=channels * nsrc, out_channels=32,
                                                 downsample=(2, 2), activation=activation, momentum=momentum)
        self.encoder_block2 = EncoderBlockRes4B(in_channels=32, out_channels=64,
                                                 downsample=(2, 2), activation=activation, momentum=momentum)
        self.encoder_block3 = EncoderBlockRes4B(in_channels=64, out_channels=128,
                                                 downsample=(2, 2), activation=activation, momentum=momentum)
        self.encoder_block4 = EncoderBlockRes4B(in_channels=128, out_channels=256,
                                                 downsample=(2, 2), activation=activation, momentum=momentum)
        self.encoder_block5 = EncoderBlockRes4B(in_channels=256, out_channels=384,
                                                 downsample=(2, 2), activation=activation, momentum=momentum)
        self.encoder_block6 = EncoderBlockRes4B(in_channels=384, out_channels=384,
                                                 downsample=(2, 2), activation=activation, momentum=momentum)
        self.conv_block7 = ConvBlockRes(in_channels=384, out_channels=384,
                                        kernel_size=(3, 3), activation=activation, momentum=momentum)
        self.decoder_block1 = DecoderBlockRes4B(in_channels=384, out_channels=384,
                                                 stride=(2, 2), activation=activation, momentum=momentum)
        self.decoder_block2 = DecoderBlockRes4B(in_channels=384, out_channels=384,
                                                 stride=(2, 2), activation=activation, momentum=momentum)
        self.decoder_block3 = DecoderBlockRes4B(in_channels=384, out_channels=256,
                                                 stride=(2, 2), activation=activation, momentum=momentum)
        self.decoder_block4 = DecoderBlockRes4B(in_channels=256, out_channels=128,
                                                 stride=(2, 2), activation=activation, momentum=momentum)
        self.decoder_block5 = DecoderBlockRes4B(in_channels=128, out_channels=64,
                                                 stride=(2, 2), activation=activation, momentum=momentum)
        self.decoder_block6 = DecoderBlockRes4B(in_channels=64, out_channels=32,
                                                 stride=(2, 2), activation=activation, momentum=momentum)

        self.after_conv_block1 = ConvBlockRes(in_channels=32, out_channels=32,
                                                 kernel_size=(3,3), activation=activation, momentum=momentum)

        self.after_conv2 = nn.Conv2d(in_channels=32, out_channels=1,
                                     kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=True)

        # self.conv = nn.Conv1d(in_channels=channels,out_channels=1,kernel_size=1)

        self.init_weights()

    def init_weights(self):
        init_layer(self.after_conv2)

    def forward(self, sp, wav):
        """
        Args:
          input: (batch_size, channels_num, segment_samples)

        Outputs:
          output_dict: {
            'wav': (batch_size, channels_num, segment_samples),
            'sp': (batch_size, channels_num, time_steps, freq_bins)}
        """
        _, cos_in, sin_in = self.f_helper.wav_to_spectrogram_phase(wav)
        # shapes: (batch_size, channels_num, time_steps, freq_bins)

        x = sp
        # (batch_size, chanenls, time_steps, freq_bins)

        # Pad spectrogram to be evenly divided by downsample ratio.
        origin_len = x.shape[2]
        pad_len = int(np.ceil(x.shape[2] / self.time_downsample_ratio)) \
                  * self.time_downsample_ratio - origin_len
        x = F.pad(x, pad=(0, 0, 0, pad_len))
        # (batch_size, channels, padded_time_steps, freq_bins)

        # Let frequency bins be evenly divided by 2, e.g., 1025 -> 1024.
        x = x[..., 0: x.shape[-1] - 1]  # (bs, channels, T, F)

        (x1_pool, x1) = self.encoder_block1(x)  # x1_pool: (bs, 32, T / 2, F / 2)
        (x2_pool, x2) = self.encoder_block2(x1_pool)  # x2_pool: (bs, 64, T / 4, F / 4)
        (x3_pool, x3) = self.encoder_block3(x2_pool)  # x3_pool: (bs, 128, T / 8, F / 8)
        (x4_pool, x4) = self.encoder_block4(x3_pool)  # x4_pool: (bs, 256, T / 16, F / 16)
        (x5_pool, x5) = self.encoder_block5(x4_pool)  # x5_pool: (bs, 512, T / 32, F / 32)
        (x6_pool, x6) = self.encoder_block6(x5_pool)  # x6_pool: (bs, 1024, T / 64, F / 64)
        x_center = self.conv_block7(x6_pool)  # (bs, 2048, T / 64, F / 64)
        x7 = self.decoder_block1(x_center, x6, both=True)  # (bs, 1024, T / 32, F / 32)
        x8 = self.decoder_block2(x7, x5, both=True)  # (bs, 512, T / 16, F / 16)
        x9 = self.decoder_block3(x8, x4, both=True)  # (bs, 256, T / 8, F / 8)
        x10 = self.decoder_block4(x9, x3, both=True)  # (bs, 128, T / 4, F / 4)
        x11 = self.decoder_block5(x10, x2, both=True)  # (bs, 64, T / 2, F / 2)
        x12 = self.decoder_block6(x11, x1, both=True)  # (bs, 32, T, F)
        x = self.after_conv_block1(x12)  # (bs, 32, T, F)
        x = self.after_conv2(x)  # (bs, channels, T, F)

        # Recover shape
        x = F.pad(x, pad=(0, 1))  # Pad frequency, e.g., 1024 -> 1025.
        x = x[:, :, 0: origin_len, :]  # (bs, channels * 3, T, F)

        mask_mag = x[:,0:1, :, :]

        out_mag = mask_mag

        out_real = out_mag * cos_in
        out_imag = out_mag * sin_in

        length = wav.shape[2]

        wav_out = self.f_helper.istft(out_real, out_imag, length)
        output_dict = {'wav': wav_out[:, None, :]}

        # wav_out = self.conv(wav)
        #
        # output_dict = {'wav': wav_out}

        return output_dict

if __name__ == "__main__":
    fh = FDomainHelper()
    model = UNetResComplex_100Mb(channels=1)
    wav = torch.randn((1,1,44100))
    f = to_log(fh.wav_to_spectrogram(wav))
    out = model(f,wav)['wav']