import math
import cairo

from typing import Callable

from xmet.sounding import Sounding

PRESSURE_MAX  = 1050 # millibar
PRESSURE_MIN  =  100
PRESSURE_STEP =   50

PRESSURE_LOG_MAX   = math.log(PRESSURE_MAX)
PRESSURE_LOG_MIN   = math.log(PRESSURE_MIN)
PRESSURE_LOG_RANGE = PRESSURE_LOG_MAX - PRESSURE_LOG_MIN

TEMP_MAX        =  50 # degrees C
TEMP_MIN        = -50
TEMP_RANGE      = TEMP_MAX - TEMP_MIN
TEMP_CENTER     =   0 # degrees C
TEMP_STEP       =   5
TEMP_STEP_COUNT = math.ceil(TEMP_RANGE / TEMP_STEP)

LAPSE_RATE_DRY   = 9.8 / 1000 # degrees C per 1000m
LAPSE_RATE_MOIST = 4.0 / 1000

def clamp(value, lowest, highest):
    if value < lowest:
        return lowest
    elif value > highest:
        return highest

    return value

class SkewTGraph():
    __slots__ = 'width', 'height', 'temp_step_width',

    def __init__(self, width: float, height: float):
        self.width  = width
        self.height = height

        self.temp_step_width = min(self.width, self.height) / TEMP_STEP_COUNT

    def graph_to_screen(self, x, y) -> tuple:
        return (self.width / 2) + x, self.height - y

    def pressure_to_y(self, pressure: float) -> float:
        log_p  = math.log(clamp(pressure, PRESSURE_MIN, PRESSURE_MAX))
        factor = (PRESSURE_LOG_MAX - log_p) / PRESSURE_LOG_RANGE

        return factor * self.height

    def draw_isobars(self, cr: cairo.Context, x: float, y: float):
        cr.save()

        for pressure in range(PRESSURE_MIN, PRESSURE_MAX+1, PRESSURE_STEP):
            coords = self.graph_to_screen(-self.width / 2,
                                         self.pressure_to_y(pressure))

            if pressure % (2*PRESSURE_STEP) == 0:
                cr.set_source_rgb(0.35, 0.35, 0.35)
            else:
                cr.set_source_rgb(0.75, 0.75, 0.75)

            cr.move_to(x + coords[0], y + coords[1])
            cr.rel_line_to(self.width, 0)
            cr.stroke()

        cr.restore()

    def skew_t_to_graph(self, x: float, y: float):
        return (x+y, y)

    def sample_to_graph(self, temp: float, pressure: float):
        x = (temp / TEMP_STEP) * self.temp_step_width
        y = self.pressure_to_y(pressure)

        return self.skew_t_to_graph(x, y)

    def sample_to_screen(self, temp: float, pressure: float):
        return self.graph_to_screen(*self.sample_to_graph(temp, pressure))

    def draw_isotherms(self, cr: cairo.Context, x: float, y: float):
        cr.save()

        cr.set_source_rgb(0.5, 0.5, 0.5)
        cr.set_line_width(0.5)
        cr.set_dash([5, 5], 1)

        for temp in range(-150, TEMP_MAX+1, TEMP_STEP):
            x1, y1 = self.sample_to_screen(temp, PRESSURE_MAX)
            x2, y2 = self.sample_to_screen(temp, PRESSURE_MIN)

            if temp % (2*TEMP_STEP) == 0:
                cr.set_source_rgb(0.35, 0.35, 0.35)
            else:
                cr.set_source_rgb(0.75, 0.75, 0.75)

            cr.move_to(x + x1, y + y1)
            cr.line_to(x + x2, y + y2)
            cr.stroke()

        cr.restore()

    def draw_sounding(self,
                     cr: cairo.Context,
                     x: float,
                     y: float,
                     sounding: Sounding,
                     fn: Callable):
        cr.save()

        first = True

        for sample in sounding.samples:
            if sample.pressure < 0 or sample.pressure is None:
                continue

            if sample.pressure < PRESSURE_MIN:
                break

            #
            # Temperature may possibly be dewpoint, depending on the
            # return value of the callback.
            #
            temp = fn(sample)

            if temp is None:
                continue

            gx, gy = self.sample_to_graph(temp, sample.pressure)
            sx, sy = self.graph_to_screen(gx, gy)

            if first:
                cr.move_to(x + sx, y + sy)
                first = False
            else:
                cr.line_to(x + sx, y + sy)

        cr.stroke()
        cr.restore()

    def draw(self,
             cr: cairo.Context,
             x: float,
             y: float,
             sounding: Sounding):
        cr.rectangle(x, y, self.width, self.height)
        cr.clip()

        self.draw_isotherms(cr, x, y)
        self.draw_isobars(cr, x, y)

        cr.set_source_rgb(1, 0, 0)
        self.draw_sounding(cr, x, y, sounding, lambda s: s.temp)

        cr.set_source_rgb(0, 1, 0)
        self.draw_sounding(cr, x, y, sounding, lambda s: s.dewpoint)

        cr.reset_clip()

class SkewTLegend():
    @staticmethod
    def draw_isotherm_legends_for_graph(cr: cairo.Context,
                                        skew_t: SkewTGraph,
                                        x: float,
                                        y: float):
        temp = TEMP_MAX

        while True:
            text    = "%d°C" % (temp)
            extents = cr.text_extents(text)

            x_rel, y_rel = skew_t.sample_to_screen(temp, PRESSURE_MAX)

            if x_rel < 0:
                y_rel = skew_t.height + x_rel
                x_rel = 0

                if y_rel <= 0:
                    break

                cr.move_to(x + x_rel - extents.width - 8,
                           y + y_rel + extents.height / 2)
            else:
                cr.move_to(x + x_rel - extents.width / 2,
                           y + y_rel + extents.height + 8)

            cr.show_text(text)
            cr.stroke()

            temp -= 2 * TEMP_STEP

    @staticmethod
    def draw_isobar_legends_for_graph(cr: cairo.Context,
                                      skew_t: SkewTGraph,
                                      x: float,
                                      y: float):
        for pressure in range(PRESSURE_MAX-PRESSURE_STEP, PRESSURE_MIN-PRESSURE_STEP-1, -2*PRESSURE_STEP):
            x_rel = skew_t.width
            y_rel = skew_t.height - skew_t.pressure_to_y(pressure)

            text    = "%dmb" % (pressure)
            extents = cr.text_extents(text)

            cr.move_to(x + x_rel + 8,
                       y + y_rel + extents.height / 2)

            cr.show_text(text)
            cr.stroke()

    @staticmethod
    def draw_for_graph(cr: cairo.Context,
                       skew_t: SkewTGraph,
                       x: float,
                       y: float):
        SkewTLegend.draw_isotherm_legends_for_graph(cr, skew_t, x, y)
        SkewTLegend.draw_isobar_legends_for_graph(cr, skew_t, x, y)