import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors, cm
from matplotlib.collections import PatchCollection
import scipy.interpolate as interpolate
from scipy import ndimage
import math
from io import BytesIO
import pandas as pd
from shapely.geometry import Point, MultiPoint
from descartes import PolygonPatch

parameters = {
    'width': {
        'default': 900,
        'type': 'int',
    },
    'height': {
        'default': 450,
        'type': 'int',
    },
    'mountain_ratio': {
        'default': 0.3,
        'type': 'float',
        'min': 0,
        'max': 1,
        'step': 0.01
    },
    'sharpness': {
        'default': 0.7,
        'type': 'float',
        'min': 0,
        'max': 1,
        'step': 0.01
    },
    'max_elevation': {
        'default': 30,
        'type': 'int',
        'min': 0,
        'max': 50,
    },
    'ground_noise': {
        'default': 15,
        'type': 'int',
        'min': 0,
        'max': 50,
    },
    'water_proportion': {
        'default': 0.6,
        'type': 'float',
        'min': 0,
        'max': 0.99,
        'step': 0.01
    },
    'mountain_concentration': {
        'default': 1,
        'type': 'float',
        'min': 0,
        'max': 5,
        'step': 0.1
    },
    'mountain_sea_distance': {
        'default': 50,
        'type': 'int',
        'min': 0,
        'max': 200,
    },
    'mountain_sea_threshold': {
        'default': 2,
        'type': 'int',
        'min': 0,
        'max': 5,
    },
    'water_level': {
        'default': 0,
        'type': 'int',
    },
    'mountain_area_elevation': {
        'default': 0.4,
        'type': 'float',
        'min': 0,
        'max': 1,
        'step': 0.01
    },
    'mountain_area_elevation_points': {
        'default': 5,
        'type': 'int',
        'min': 0,
        'max': 15,
    },
    'mountain_area_elevation_area': {
        'default': 10,
        'type': 'int',
        'min': 0,
        'max': 25,
    },
    'continents': {
        'default': 5,
        'type': 'int',
    },
    'continent_spacing': {
        'default': 0.3,
        'type': 'float',
        'min': 0,
        'max': 1,
        'step': 0.1
    },
    'biomes': {
        'default': False,
        'type': 'bool'
    },
    'seed': {
        'default': '',
        'type': 'int',
        'description': 'Leave empty for a random seed generated from the current timestamp.'
    },
        
}
p = { k: parameters[k]['default'] for k in parameters }

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

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 > p['water_level']

# TODO: should check as a sphere
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

    if x < 0:
        x = w + x
    elif x >= w:
        x = x - w

    if y < 0:
        y = h + y
    elif y >= h:
        y = y - h

    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 trials > CONTINENT_MAX_TRIALS:
            # print('couldnt proceed')
        if size <= 0 or trials > CONTINENT_MAX_TRIALS: break
        # if size <= 0: 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

        x, y = bound_check(ground, (x, y))

        if not is_ground(ground[x, y]) and in_range((x, y), position, size**2 * np.pi):
            trials = 0
            size -= 1
            ground[x, y] = np.random.randint(1, p['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=p['mountain_sea_distance']):
    ns = neighbours(ground, position, radius).flatten()
    sea = len([1 for x in ns if not is_ground(x)])

    return sea < p['mountain_sea_threshold']

def random_elevate_agent(ground, position, height, size=p['mountain_area_elevation_points']):
    position = position + np.random.random_integers(-p['mountain_area_elevation_area'], p['mountain_area_elevation_area'], size=2)

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

        change = height * p['mountain_area_elevation']
        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(p['max_elevation'])

    ground[x, y] = height

    last_height = height
    for i in range(1, height):
        for d in DIRECTIONS:
            change = np.random.randint(p['mountain_concentration'] + 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 - change
        if last_height < 0:
            break

    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 - p['sharpness']):
        return max(1, a[0])
    return 0

def generate_map(biomes=False, **kwargs):
    plt.clf()

    p.update(kwargs)

    np.random.seed(p['seed'] or None)

    width, height = p['width'], p['height']
    continents = p['continents']

    ground = np.zeros((width, height))
    ground_size = width * height * (1 - p['water_proportion'])
    print(ground_size / ground.size)

    # position = (int(width / 2), int(height / 2))
    # ground_size = width * height * GROUND_PROPORTION
    # continent_agent(ground, position, size=ground_size)
    position = (0, int(height / 2))
    ym = 1
    for continent in range(continents):
        position = (position[0] + np.random.randint(p['continent_spacing'] * width * 0.8, p['continent_spacing'] * width * 1.2),
                    position[1] + ym * np.random.randint(p['continent_spacing'] * height * 0.8, p['continent_spacing'] * height * 1.2))

        print(position)

        ym = ym * -1

        random_size = ground_size / continents
        continent_agent(ground, position, size=random_size)

    ground = ndimage.gaussian_filter(ground, sigma=(1 - p['sharpness']) * 20)

    for i in range(int(ground_size * p['mountain_ratio'] / p['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), p['max_elevation'])
    print(np.unique(ground))

    print(np.count_nonzero(ground) / ground.size)

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

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

    if biomes:
        generate_biomes(ground)

    return figfile

def to_range(omin, omax, nmin, nmax):
    orange = omax - omin
    nrange = nmax - nmin

    return lambda x: ((x - omin) * nrange / orange) + nmin
        

def generate_biomes(ground):
    width, height = p['width'], p['height']

    height_to_latitude = to_range(0, height, -90, 90)
    width_to_longitude = to_range(0, width, -180, 180)

    print('generate_biomes')
    INPUTS = ['elevation', 'distance_to_water', 'latitude']
    
    data = {}
    for col in ['longitude', 'latitude', 'elevation', 'distance_to_water']:
        data[col] = []

    points = []
    for x, y in np.ndindex(ground.shape):
        v = ground[x,y]
        if v > p['water_level']:
            points.append((x, y))

            data['longitude'].append(width_to_longitude(x))
            data['latitude'].append(height_to_latitude(y))
            data['elevation'].append(v)

    points = MultiPoint(points)
    boundary = points.buffer(1e-0).boundary

    for x, y in np.ndindex(ground.shape):
        if ground[x,y] > p['water_level']:
            # print(x,y, Point(x,y).distance(boundary))
            data['distance_to_water'].append(Point(x, y).distance(boundary))

    
    df = pd.DataFrame(data)
    print(df)

    # fig = plt.figure()
    # ax = fig.add_subplot(111)
    # minx, miny, maxx, maxy = boundary.bounds
    # w, h = maxx - minx, maxy - miny
    # ax.set_xlim(minx - 0.2 * w, maxx + 0.2 * w)
    # ax.set_ylim(miny - 0.2 * h, maxy + 0.2 * h)
    # ax.set_aspect(1)

    # ax.add_collection(PatchCollection([PolygonPatch(boundary.buffer(0.1), fc='red', ec='black', zorder=1)], match_original=True))

    # plt.show()


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