import numpy as np
import sys
import time
import matplotlib
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

plt.rc('font', family='serif',size=24)
matplotlib.rc('text', usetex=True)
matplotlib.rc('legend', fontsize=24)
matplotlib.rcParams['text.latex.preamble'] = r'\boldmath'


# Observation times.
# X should be in units of minutes since midnight
# Y should be in units of ppb
OBS_X = [456, 486, 600, 690, 780]
OBS_Y = [2600, 350, 200, 100, 60]
OBS_Y = [x/2 for x in OBS_Y]

# These are the twilight boundaries for a particular day.  The order is:
# 1. Midnight
# 2. Astronomical Twilight begins
# 3. Nautical Twilight begins
# 4. Civil Twilight begins
# 5. Daylight begins
# 6. Civil Daylight ends
# 7. Nautical Daylight ends
# 8. Astronomical Daylight ends
# 9. Astronmical Twilight ends
# 10. Midnight - 1
# Units are in minutes since midnight
# Currently set for March 28th, 2024 in Pampa, TX
TWI_TIMES = [0, 369, 399, 429, 455, 1202, 1228, 1258, 1289, 1439]

COLOR_BLACK = (0,0,0)
COLOR_DARKGRAY = (63,63,63)
COLOR_GRAY = (127,127,127)
COLOR_LIGHTGRAY = (191,191,191)
COLOR_WHITE = (255,255,255)
COLOR_SKY = (106, 153, 166)
COLOR_SKYD = (79, 121, 133)
COLOR_SKYDD = (57, 87, 95)
COLOR_SKYDDD = (35, 53, 58)

class Point(object):
    def __init__(self, x, y):
        self.x, self.y = x, y

class Rect(object):
    def __init__(self, x1, y1, x2, y2):
        minx, maxx = (x1,x2) if x1 < x2 else (x2,x1)
        miny, maxy = (y1,y2) if y1 < y2 else (y2,y1)
        self.min = Point(minx, miny)
        self.max = Point(maxx, maxy)

    width  = property(lambda self: self.max.x - self.min.x)
    height = property(lambda self: self.max.y - self.min.y)


def gradient_color(minval, maxval, val, color_palette):
    """ Computes intermediate RGB color of a value in the range of minval
        to maxval (inclusive) based on a color_palette representing the range.
    """
    max_index = len(color_palette)-1
    delta = maxval - minval
    if delta == 0:
        delta = 1
    v = float(val-minval) / delta * max_index
    i1, i2 = int(v), min(int(v)+1, max_index)
    (r1, g1, b1), (r2, g2, b2) = color_palette[i1], color_palette[i2]
    f = v - i1
    return int(r1 + f*(r2-r1)), int(g1 + f*(g2-g1)), int(b1 + f*(b2-b1))

def horz_gradient(draw, rect, color_func, color_palette):
    minval, maxval = 1, len(color_palette)
    delta = maxval - minval
    width = float(rect.width)  # Cache.
    for x in range(rect.min.x, rect.max.x+1):
        f = (x - rect.min.x) / width
        val = minval + f * delta
        color = color_func(minval, maxval, val, color_palette)
        draw.line([(x, rect.min.y), (x, rect.max.y)], fill=color)

#def vert_gradient(draw, rect, color_func, color_palette):
    #minval, maxval = 1, len(color_palette)
    #delta = maxval - minval
    #height = float(rect.height)  # Cache.
    #for y in range(rect.min.y, rect.max.y+1):
        #f = (y - rect.min.y) / height
        #val = minval + f * delta
        #color = color_func(minval, maxval, val, color_palette)
        #draw.line([(rect.min.x, y), (rect.max.x, y)], fill=color)

def rise_assmb(in_times):
    [start, astwi1, natwi1, citwi1, day1, day2, citwi2, natwi2, astwi2, end] = in_times
    #night = [COLOR_BLACK, COLOR_BLACK]
    #astro = [COLOR_BLACK, COLOR_DARKGRAY]
    #nauti = [COLOR_DARKGRAY, COLOR_GRAY]
    #civil = [COLOR_LIGHTGRAY, COLOR_WHITE]
    #white = [COLOR_WHITE, COLOR_WHITE]
    night = [COLOR_BLACK, COLOR_BLACK]
    astro = [COLOR_BLACK, COLOR_SKYDDD]
    nauti = [COLOR_SKYDDD, COLOR_SKYDD]
    civil = [COLOR_SKYDD, COLOR_SKYD]
    white = [COLOR_SKY, COLOR_SKY]
    colorgrad = [night, astro, nauti, civil, white, civil[::-1], nauti[::-1], astro[::-1], night]
    times = [start, astwi1, natwi1, citwi1, day1, day2, citwi2, natwi2, citwi2, end]
    allimg = None
    for i, block in enumerate(colorgrad):
        time1 = times[i]
        time2 = times[i+1]
        grad = colorgrad[i]
        region = Rect(0, 0, time2-time1, 1)
        width, height = region.max.x+1, region.max.y+1
        localimg = Image.new("RGB", (width, height), COLOR_WHITE)
        draw = ImageDraw.Draw(localimg)
        horz_gradient(draw, region, gradient_color, colorgrad[i])
        if allimg == None:
            allimg = localimg
        else:
            imgs = [allimg, localimg]
            allimg = Image.fromarray(
                np.concatenate(
                    [np.array(x) for x in imgs],
                    axis=1
                )
            )
    return allimg



if __name__ == '__main__':

    allimg = rise_assmb(TWI_TIMES)
    #allimg.show()
    #time.sleep(1)

    fig = plt.figure(figsize=[16,16])
    ax = fig.add_subplot(111)
    #plt.plot([1,2,3],[1,2,1])
    plotlim = plt.xlim() + plt.ylim()
    ax.imshow(allimg, extent=[0, 1440, 0, 1500])
    #plt.xlim(300, 1000)
    plt.scatter(OBS_X, OBS_Y, s=120, c='red')
    lab = ax.axes.get_xticks().tolist()

    ax.yaxis.set_ticks(np.arange(0, 1500, 250))
    tlab = []
    for i in range(0, 1439):
        if i % 250 == 0:
            tlab.append(str(int(i*2)))
    ax.set_yticklabels(tlab)

    ax.xaxis.set_ticks(np.arange(0, 1440, 60))
    tlab = []
    for i in range(0, 1439):
        if i % 60 == 0:
            tlab.append(str(int(i/60)).zfill(2))
    ax.set_xticklabels(tlab)

    ax.set_xlabel("Hour of day, local timezone")
    ax.set_ylabel("CH$_4$ enhancement (ppb)")


    plt.savefig("./daylight_ch4_enhancement.png")