import os
import re
import random
import imageHandling

from collections import namedtuple
from typing import Any, List, Dict

from imageHandling import Color as Color
from imageHandling import Image as Image

import pprint
pp = pprint.PrettyPrinter(indent=4)


CREST_WIDTH = 17
CREST_HEIGHT = 20

BG_COLOR = Color(124, 139, 150, 255)


Vector = namedtuple('Vector', 'x, y')


class ColorShades:
    def __init__(self, light: Color, medium: Color, dark: Color):
        self.light = light
        self.medium = medium
        self.dark = dark


ColorSet = [ColorShades, ColorShades, ColorShades, ColorShades]


class ShieldData:
    def __init__(self, color: Image, shade: Image, charge: Image):
        self.color = color
        self.shade = shade
        self.charge = charge


class ChargeData:
    def __init__(self, color: Image, shade: Image):
        self.color = color
        self.shade = shade


class ChargeSpaceData:
    def __init__(self, start: Vector, end: Vector, color: Color):
        self.start = start
        self.end = end
        self.color = color


class ShadedColor:
    def __init__(self, color_shades: ColorShades, shade_value: Color, is_empty: bool):
        self.color_shades = color_shades
        self.shade_value = shade_value
        self.is_empty = is_empty


class ShadeableImage:
    def __init__(self, width: int, height: int):
        self.width = width
        self.height = height
        default_value = ShadedColor(ColorShades(Color.clear(), Color.clear(), Color.clear()), Color.clear(), True)
        self.data = [[default_value for _ in range(width)] for _ in range(height)]

    def __str__(self):
        return "ShadeableImage(%s, %s)" % (self.width, self.height)

    def __repr__(self):
        return self.__str__()

    def get(self, x: int, y: int) -> ShadedColor:
        return self.data[y][x]

    def set(self, x: int, y: int, shaded_color: ShadedColor) -> 'ShadeableImage':
        if type(shaded_color) != ShadedColor:
            raise ValueError('shaded_color is not of type ShadedColor')
        self.data[y][x] = shaded_color
        return self

    def cut_out(self, offset_x: int, offset_y: int, width: int, height: int) -> 'ShadeableImage':
        if offset_x < 0 or offset_x >= self.width:
            raise ValueError("start_x is out of bounds")
        if offset_y < 0 or offset_y >= self.height:
            raise ValueError("start_y is out of bounds")
        result = ShadeableImage(width, height)
        for y in range(height):
            for x in range(width):
                result.set(x, y, self.get(x + offset_x, y + offset_y))
        return result

    def set_range(self, offset_x: int, offset_y: int, image: 'ShadeableImage',
                  operation=lambda x, y: y) -> 'ShadeableImage':
        if offset_x < 0 or offset_x >= self.width:
            raise ValueError("start_x is out of bounds")
        if offset_y < 0 or offset_y >= self.height:
            raise ValueError("start_y is out of bounds")
        for y in range(image.height):
            for x in range(image.width):
                self.set(x + offset_x, y + offset_y, operation(self.get(x, y), image.get(x, y)))
        return self

    def fill(self, shaded_color: ShadedColor) -> 'ShadeableImage':
        for y in range(self.height):
            for x in range(self.width):
                self.set(x, y, shaded_color)
        return self

    def flip_horizontal(self) -> 'ShadeableImage':
        image = ShadeableImage(self.width, self.height)
        for x in range(self.width):
            for y in range(self.height):
                image.set(x, y, self.get(self.width - x - 1, y))
        self.data = image.data
        return self

    def to_image(self) -> Image:
        result = Image(self.width, self.height)
        for y in range(self.height):
            for x in range(self.width):
                data = self.get(x, y)
                if data.is_empty:
                    result.set(x, y, BG_COLOR)
                else:
                    result.set(x, y, get_color_from_shade(data.color_shades, data.shade_value))
        return result


COLORS = [
    [Color(189, 0, 0, 255), Color(166, 4, 4, 255), Color(141, 3, 3, 255)],  # red
    [Color(0, 75, 25, 255), Color(0, 64, 30, 255), Color(0, 46, 28, 255)],  # green
    [Color(0, 54, 135, 255), Color(0, 35, 105, 255), Color(0, 28, 84, 255)],  # blue
    [Color(0, 153, 213, 255), Color(18, 131, 195, 255), Color(15, 100, 164, 255)],  # light blue
    [Color(28, 28, 28, 255), Color(14, 14, 14, 255), Color(0, 0, 0, 255)],  # black
    # [Color(76, 26, 166, 255), Color(64, 18, 128, 255), Color(43, 15, 110, 255)]  # purple
]

METALS = [
    [Color(225, 175, 0, 255), Color(209, 150, 0, 255), Color(184, 113, 0, 255)],  # yellow
    [Color(255, 255, 255, 255), Color(227, 227, 227, 255), Color(196, 196, 196, 255)],  # white
]

CHARGE_SIZES = [
    Vector(3, 3)
]

PROTOTYPES = [
    "bottom_s",
    "bottom_w",
    "cantonlarge",
    "cantonsmall",
    "chieftri_s",
    "chieftri_w",
    "diagonal",
    "horizontal_s",
    "horizontal_w",
    "side",
    "top_s",
    "top_w",
    "vertical"
]


def read_tinctures_from_file(filename: str) -> List[ColorShades]:
    image = imageHandling.image_from_file(filename)
    result = []
    for y in range(image.height):
        result.append(ColorShades(image.get(0, y), image.get(1, y), image.get(2, y)))
    return result


def read_charge_sizes_from_file(directory: str) -> List[Vector]:
    result = []
    for filename in os.listdir(os.path.abspath(directory)):
        match = re.match("([0-9]+)x([0-9]+).png", filename)
        if match:
            result.append(Vector(int(match.group(1)), int(match.group(2))))
    return result


def random_shield_data_from_file(filename: str) -> ShieldData:
    image = imageHandling.image_from_file(filename)
    offset_x = int(random.random() * image.width / CREST_WIDTH) * CREST_WIDTH
    color = image.cut_out(offset_x, 0, CREST_WIDTH, CREST_HEIGHT)
    shade = image.cut_out(offset_x, CREST_HEIGHT, CREST_WIDTH, CREST_HEIGHT)
    charge = image.cut_out(offset_x, 2 * CREST_HEIGHT, CREST_WIDTH, CREST_HEIGHT)
    return ShieldData(color, shade, charge)


def random_charge_data_from_file(filename: str, charge_width: int, charge_height) -> ChargeData:
    image = imageHandling.image_from_file(filename)
    offset_x = int(random.random() * image.width / charge_width) * charge_width
    color = image.cut_out(offset_x, 0, charge_width, charge_height)
    shade = image.cut_out(offset_x, charge_height, charge_width, charge_height)
    return ChargeData(color, shade)


def find_charge_spaces(charge_data: Image) -> List[ChargeSpaceData]:
    result = []
    for y in range(1, charge_data.height):
        for x in range(1, charge_data.width):
            if ((not charge_data.get(x, y).is_transparent()) and
                    charge_data.get(x - 1, y).is_transparent() and
                    charge_data.get(x, y - 1).is_transparent()):
                end = find_charge_space_end(charge_data, Vector(x, y))
                result.append(ChargeSpaceData(Vector(x, y), end, charge_data.get(x, y)))
    return result


def find_charge_space_end(charge_data: Image, start: Vector) -> Vector:
    result_x = result_y = 0
    base_color = charge_data.get(start.x, start.y)
    for x in range(start.x + 1, charge_data.width):
        if charge_data.get(x, start.y) != base_color:
            result_x = x - 1
            break
    for y in range(start.y + 1, charge_data.height):
        if charge_data.get(start.x, y) != base_color:
            result_y = y - 1
            break
    return Vector(result_x, result_y)


def combine_colors(a: Color, b: Color) -> Color:
    if b.is_transparent():
        return a
    return b


def random_color_set():
    if random.random() > 0.5:
        a = random.sample(COLORS, 2)
        b = random.sample(METALS, 2)
    else:
        a = random.sample(METALS, 2)
        b = random.sample(COLORS, 2)
    return a[0], b[0], a[1], b[1]


def calculate_color(color_set: ColorSet, color_data: Color, alt_color: [bool, bool]) -> ColorShades:
    if color_data.is_red() and not color_data.is_transparent():
        return color_set[0]
    if color_data.is_blue() and not color_data.is_transparent():
        return color_set[1]
    if color_data.is_magenta() and not color_data.is_transparent():
        if alt_color[0]:
            return color_set[2]
        else:
            return color_set[0]
    if color_data.is_cyan() and not color_data.is_transparent():
        if alt_color[1]:
            return color_set[3]
        else:
            return color_set[1]
    raise ValueError('unknown color code')


def get_color_from_shade(color_shades: ColorShades, shade_data: Color) -> Color:
    if shade_data.is_white() and not shade_data.is_transparent():
        return color_shades.light
    if shade_data.is_gray() and not shade_data.is_transparent():
        return color_shades.medium
    if shade_data.is_black() and not shade_data.is_transparent():
        return color_shades.dark
    return Color.magenta()


def random_shield() -> Image:
    prototype = random.choice(PROTOTYPES)
    file_a = os.path.abspath("image/prototypes/" + prototype + "_a.png")
    file_b = os.path.abspath("image/prototypes/" + prototype + "_b.png")
    color_set = random_color_set()

    part_a = random_shield_data_from_file(file_a)
    part_b = random_shield_data_from_file(file_b)
    image = ShadeableImage(CREST_WIDTH, CREST_HEIGHT)
    paint_part_to_image(image, part_a, color_set)
    paint_part_to_image(image, part_b, list(reversed(color_set)))

    return image.to_image()


def paint_part_to_image(image: ShadeableImage, data: ShieldData, color_set: ColorSet):
    alt_color = [random.random() > 0.5, random.random() > 0.5]
    for x in range(CREST_WIDTH):
        for y in range(CREST_HEIGHT):
            color_data = data.color.get(x, y)
            shade_data = data.shade.get(x, y)
            if color_data.is_transparent():
                continue
            color_shades = calculate_color(color_set, color_data, alt_color)
            image.set(x, y, ShadedColor(color_shades, shade_data, False))

    # add charge
    charge_spaces = find_charge_spaces(data.charge)
    charge_space_dictionary = {}
    charge_dictionary = {}
    for charge_space in charge_spaces:
        adjusted_charge_spaces = adjust_charge_space(charge_space, charge_space_dictionary)
        alt_color = [random.random() > 0.5, random.random() > 0.5]
        for adjusted_charge_space in adjusted_charge_spaces:
            add_charge_to_image(adjusted_charge_space, charge_dictionary, image, color_set, alt_color, data)


def adjust_charge_space(charge_space: ChargeSpaceData,
                        charge_space_dictionary: Dict[Any, ChargeSpaceData]) -> List[ChargeSpaceData]:
    rnd = random.random() * 255
    if rnd > charge_space.color.a:
        return []
    charge_space_width = charge_space.end.x - charge_space.start.x + 1
    charge_space_height = charge_space.end.y - charge_space.start.y + 1
    charge_space_dictionary_key = (charge_space_width, charge_space_height, charge_space.color)

    # select charge space combination
    if charge_space_dictionary_key in charge_space_dictionary:
        adjusted_charge_spaces = charge_space_dictionary[charge_space_dictionary_key]
    else:
        adjusted_charge_spaces_list = convert_charge_space(charge_space)
        if len(adjusted_charge_spaces_list) == 0:
            return []
        adjusted_charge_spaces = random.choice(adjusted_charge_spaces_list)
        charge_space_dictionary[charge_space_dictionary_key] = adjusted_charge_spaces

    return adjusted_charge_spaces


def add_charge_to_image(charge_space: ChargeSpaceData, charge_dictionary: Dict[Any, ChargeData],
                        image: ShadeableImage, color_set: ColorSet, alt_color: (bool, bool), data: ShieldData) -> None:
    charge_space_width = charge_space.end.x - charge_space.start.x + 1
    charge_space_height = charge_space.end.y - charge_space.start.y + 1
    charge_dictionary_key = (charge_space_width, charge_space_height, charge_space.color)
    if charge_dictionary_key in charge_dictionary:
        charge_data = charge_dictionary[charge_dictionary_key]
    else:
        charge_data = get_charge_for_size(charge_space_width, charge_space_height)
        charge_dictionary[charge_dictionary_key] = charge_data
    for x in range(charge_data.color.width):
        for y in range(charge_data.color.height):
            image_x = charge_space.start.x + x
            image_y = charge_space.start.y + y
            shade = combine_shade_code(data.shade.get(image_x, image_y), charge_data.shade.get(x, y))
            if charge_data.color.get(x, y).is_white():
                color = image.get(image_x, image_y).color_shades
            else:
                color = calculate_color(color_set, charge_space.color, alt_color)
            shaded_color = ShadedColor(color, shade, False)
            image.set(image_x, image_y, shaded_color)


def combine_shade_code(a: Color, b: Color) -> Color:
    if a == Color.black() or b == Color.black():
        return Color.black()
    elif a == b == Color.gray():
        return Color.black()
    elif a == Color.gray() or b == Color.gray():
        return Color.gray()
    else:
        return Color.white()


def convert_charge_space(original: ChargeSpaceData) -> List[List[ChargeSpaceData]]:
    result = []
    for charge_size in CHARGE_SIZES:
        original_size, difference, factor = calculate_charge_space_split_data(original, charge_size)
        # sort out larger charges
        if charge_size.x > original_size.x or charge_size.y > original_size.y:
            continue
        # add same size charge
        if charge_size.x == original_size.x and charge_size.y == original_size.y:
            result.append([original])
            continue
        # add smaller single charge
        difference = Vector(original_size.x - charge_size.x, original_size.y - charge_size.y)
        if difference.x % 2 == 0:
            new_start = Vector(original.start.x + difference.x // 2, original.start.y + difference.y // 2)
            new_end = Vector(new_start.x + charge_size.x - 1, new_start.y + charge_size.y - 1)
            result.append([ChargeSpaceData(new_start, new_end, original.color)])
        # add multiple
        for combination in get_horizontal_line_charges(original, charge_size):
            result.append(combination)
        for combination in get_vertical_line_charges(original, charge_size):
            result.append(combination)
        for combination in get_large_combination_charges(original, charge_size):
            result.append(combination)
    return result


def get_horizontal_line_charges(original: ChargeSpaceData, charge_size: Vector) -> List[List[ChargeSpaceData]]:
    result = []
    original_size, difference, factor = calculate_charge_space_split_data(original, charge_size)
    if factor.x <= 1:
        return result
    new_start_y = original.start.y + difference.y // 2
    for i in range(2, factor.x + 1):
        rest = original_size.x - charge_size.x * i
        if rest % (i - 1) != 0 or rest < i - 1:
            continue
        gutter = rest // (i - 1)
        new_charge_space_set = []
        for j in range(i):
            new_start = Vector(original.start.x + j * (gutter + charge_size.x), new_start_y)
            new_end = Vector(new_start.x + charge_size.x - 1, new_start.y + charge_size.y - 1)
            new_charge_space_set.append(ChargeSpaceData(new_start, new_end, original.color))
        result.append(new_charge_space_set)
    return result


def get_vertical_line_charges(original: ChargeSpaceData, charge_size: Vector) -> List[List[ChargeSpaceData]]:
    result = []
    original_size, difference, factor = calculate_charge_space_split_data(original, charge_size)
    if factor.y <= 1 or difference.x % 2 != 0:
        return result
    new_start_x = original.start.x + difference.x // 2
    for i in range(2, factor.y + 1):
        rest = original_size.y - charge_size.y * i
        if rest < i - 1:
            continue
        gutter = rest // (i - 1)
        new_charge_space_set = []
        for j in range(i):
            new_start = Vector(new_start_x, original.start.y + j * (gutter + charge_size.y))
            new_end = Vector(new_start.x + charge_size.x - 1, new_start.y + charge_size.y - 1)
            new_charge_space_set.append(ChargeSpaceData(new_start, new_end, original.color))
        result.append(new_charge_space_set)
    return result


def get_large_combination_charges(original: ChargeSpaceData, charge_size: Vector) -> List[List[ChargeSpaceData]]:
    result = []
    original_size, difference, factor = calculate_charge_space_split_data(original, charge_size)
    if original_size.x <= charge_size.x * 2 or original_size.y <= charge_size.y * 2:
        return result
    start_00 = original.start
    end_00 = Vector(start_00.x + charge_size.x - 1, start_00.y + charge_size.y - 1)
    start_11 = Vector(original.end.x - charge_size.x + 1, original.end.y - charge_size.y + 1)
    end_11 = Vector(start_11.x + charge_size.x - 1, start_11.y + charge_size.y - 1)
    start_01 = Vector(start_00.x, start_11.y)
    end_01 = Vector(end_00.x, end_11.y)
    start_10 = Vector(start_11.x, start_00.y)
    end_10 = Vector(end_11.x, end_00.y)

    result.append([
        ChargeSpaceData(start_00, end_00, original.color),
        ChargeSpaceData(start_01, end_01, original.color),
        ChargeSpaceData(start_10, end_10, original.color),
        ChargeSpaceData(start_11, end_11, original.color)
    ])
    result.append([
        ChargeSpaceData(start_00, end_00, original.color),
        ChargeSpaceData(start_11, end_11, original.color)
    ])
    result.append([
        ChargeSpaceData(start_01, end_01, original.color),
        ChargeSpaceData(start_10, end_10, original.color)
    ])

    # add with middle charge if possible
    if difference.x % 2 != 0 or (factor.x <= 2 and factor.y <= 2):
        return result
    start_middle = Vector(original.start.x + difference.x // 2, original.start.y + difference.y // 2)
    end_middle = Vector(start_middle.x + charge_size.x - 1, start_middle.y + charge_size.y - 1)
    result.append([
        ChargeSpaceData(start_00, end_00, original.color),
        ChargeSpaceData(start_01, end_01, original.color),
        ChargeSpaceData(start_10, end_10, original.color),
        ChargeSpaceData(start_11, end_11, original.color),
        ChargeSpaceData(start_middle, end_middle, original.color)
    ])
    result.append([
        ChargeSpaceData(start_00, end_00, original.color),
        ChargeSpaceData(start_11, end_11, original.color),
        ChargeSpaceData(start_middle, end_middle, original.color)
    ])
    result.append([
        ChargeSpaceData(start_01, end_01, original.color),
        ChargeSpaceData(start_10, end_10, original.color),
        ChargeSpaceData(start_middle, end_middle, original.color)
    ])
    return result


def calculate_charge_space_split_data(original: ChargeSpaceData, charge_size: Vector) -> [Vector, Vector, Vector]:
    original_size = Vector(original.end.x - original.start.x + 1, original.end.y - original.start.y + 1)
    difference = Vector(original_size.x - charge_size.x, original_size.y - charge_size.y)
    factor = Vector(original_size.x // charge_size.x, original_size.y // charge_size.y)
    return original_size, difference, factor


def get_charge_for_size(width: int, height: int) -> ChargeData:
    filename = "image/charges/%sx%s.png" % (width, height)
    if os.path.isfile(filename):
        return random_charge_data_from_file(filename, width, height)
    else:
        return ChargeData(Image(width, height).fill(Color.white()), Image(width, height).fill(Color.white()))


def random_shield_plate(width: int, height: int) -> Image:
    image = Image(width * CREST_WIDTH, height * CREST_HEIGHT)
    shield_sum = width * height
    shield_counter = 0
    for x in range(0, image.width, CREST_WIDTH):
        for y in range(0, image.height, CREST_HEIGHT):
            shield_counter += 1
            print("generate shield %s of %s" % (shield_counter, shield_sum))
            image.set_range(x, y, random_shield())
    return image


COLORS = read_tinctures_from_file("image/tinctures/colors.png")
METALS = read_tinctures_from_file("image/tinctures/metals.png")
CHARGE_SIZES = read_charge_sizes_from_file("image/charges")

imageHandling.write_image_to_file("test.png", random_shield_plate(20, 10))

# for i in range(20): imageHandling.write_image_to_file("proceduralHeraldry_%02d.png" % i, random_shield_plate(100, 50))