import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors, cm
import scipy.interpolate as interpolate
from scipy import ndimage
import math
from io import BytesIO
import base64

WIDTH = 900
HEIGHT = 450
RATIO = WIDTH / HEIGHT

MOUNTAIN_SEA_DISTANCE = 50
MOUNTAIN_SEA_THRESHOLD = 2
MOUNTAIN_JAGGEDNESS = 1

CONTINENT_MAX_TRIALS = 1e4
MOUNTAIN_RATIO = 0.3
SEA_COLOR = np.array((53, 179, 220, 255)) / 255

SHARPNESS = 0.7
WATER_LEVEL = 0
MAX_ELEVATION = 30
GROUND_NOISE = 15
WATER_PROPORTION = 0.6
GROUND_PROPORTION = 1 - WATER_PROPORTION

DIRECTIONS = [(-1, -1), (-1, 0), (-1, 1), (1, 1), (1, 0), (1, -1), (0, -1), (0, 1)]

def s(x):
    return -2 * x**3 + 3 * x**2

def is_ground(value):
    return value > WATER_LEVEL

def in_range(p, m, size):
    x, y = p
    mx, my = m
    return ((x - mx)**2 + (y - my)**2) < size

def max_recursion(fn, max_recursion=0):
    def f(*args, recursion=0, **kwargs):
        if recursion > max_recursion:
            return

        return fn(*args, **kwargs)

def bound_check(ground, point):
    x, y = point
    w, h = ground.shape

    x = max(min(x, w - 1), 0)
    y = max(min(y, h - 1), 0)

    return (x, y)


def continent_agent(ground, position, size):
    if size <= 0: return
    x, y = position

    w, h = ground.shape

    trials = 0

    while True:
        if size <= 0 or trials > CONTINENT_MAX_TRIALS: break

        dx = np.random.randint(2) or -1
        dy = np.random.randint(2) or -1

        r = np.random.randint(3)
        new_point = bound_check(ground, (x + dx, y + dy))
        if r == 0:
            x = new_point[0]
        elif r == 1:
            y = new_point[1]
        else:
            x, y = new_point


        if not is_ground(ground[x, y]) and in_range((x, y), position, size):
            trials = 0
            size -= 1
            ground[x, y] = np.random.randint(1, GROUND_NOISE)
        else:
            trials += 1

def neighbours(ground, position, radius):
    x, y = position
    return ground[x-radius:x+radius+1, y-radius:y+radius+1]

def away_from_sea(ground, position, radius=MOUNTAIN_SEA_DISTANCE):
    ns = neighbours(ground, position, radius).flatten()
    sea = len([1 for x in ns if not is_ground(x)])

    return sea < MOUNTAIN_SEA_THRESHOLD

MOUNTAIN_AREA_ELEVATION = 0.4
MOUNTAIN_AREA_ELEVATION_N = 5
MOUNTAIN_AREA_ELEVATION_AREA = 10
def random_elevate_agent(ground, position, height, size=MOUNTAIN_AREA_ELEVATION_N):
    position = position + np.random.random_integers(-MOUNTAIN_AREA_ELEVATION_AREA, MOUNTAIN_AREA_ELEVATION_AREA, size=2)

    for i in range(size):
        d = DIRECTIONS[np.random.randint(len(DIRECTIONS))]

        change = height * MOUNTAIN_AREA_ELEVATION + np.random.randint(MOUNTAIN_JAGGEDNESS + 1)
        new_index = bound_check(ground, position + np.array(d))

        if is_ground(ground[new_index]):
            ground[new_index] += change


def mountain_agent(ground, position):
    if not away_from_sea(ground, position):
        return

    x, y = position
    height = np.random.randint(MAX_ELEVATION)

    ground[x, y] = height

    last_height = height
    for i in range(1, height):
        for d in DIRECTIONS:
            change = np.random.randint(MOUNTAIN_JAGGEDNESS + 1)
            distance = np.array(d)*i
            new_index = bound_check(ground, position + distance)

            if is_ground(ground[new_index]):
                ground[new_index] = last_height - change

        last_height = last_height - MOUNTAIN_JAGGEDNESS

    random_elevate_agent(ground, position, height)


# takes an initial position and a list of (direction, probability) tuples to walk on
# def split_agent(ground, position, directions):

def constant_filter(a):
    if a[0] > (1 - SHARPNESS):
        return max(1, a[0])
    return 0

def generate_map(width=WIDTH, height=HEIGHT, continents=4):
    ground = np.zeros((width, height))
    ground_size = width * height * GROUND_PROPORTION

    # position = (int(width / 2), int(height / 2))
    # ground_size = width * height * GROUND_PROPORTION
    # continent_agent(ground, position, size=ground_size)
    for continent in range(continents):
        position = (np.random.randint(0, width), np.random.randint(0, height))
        print(position)
        continent_agent(ground, position, size=ground_size)

    ground = ndimage.gaussian_filter(ground, sigma=(1 - SHARPNESS) * 20)

    for i in range(int(ground_size * MOUNTAIN_RATIO / MAX_ELEVATION**2)):
        position = (np.random.randint(0, width), np.random.randint(0, height))
        mountain_agent(ground, position)

    norm = colors.Normalize(vmin=1)
    greys = cm.get_cmap('Greys')
    greys.set_under(color=SEA_COLOR)

    ground = ndimage.gaussian_filter(ground, sigma=4)
    ground = ndimage.generic_filter(ground, constant_filter, size=1)
    print(np.min(ground), np.max(ground), MAX_ELEVATION)
    print(np.unique(ground))

    plt.imshow(ground.T, cmap=greys, norm=norm)

    figfile = BytesIO()
    plt.savefig(figfile, format='png')
    figfile.seek(0)

    return figfile

if __name__ == "__main__":
    generate_map(WIDTH, HEIGHT)
    plt.show()