#!/usr/bin/env python

import colour
import numpy as np
from yaml import dump

'''
We want to define a 16-color palette for editors and terminals.  The palette
consists of 8 monotones and 8 colors.

For editors, the 8 monotones are for background and text.  From the darkest to
the brightest, they are

 - bg0: text bg
 - bg1: ui bg (e.g., status bar bg)
 - bg2: highlight bg (e.g., current selection)
 - fg0: ui text / comment
 - fg1: ui text highlight
 - fg2: text
 - fg3: unused
 - fg4: text highlight

The 8 colors are mainly for syntax highlighting.  We want them to differ only
in hues.  The hues are exactly 45 degrees apart from each other on the color
wheel for better contrast

 - red:     0 degrees
 - orange:  45 degrees
 - yellow:  90 degrees
 - green:   135 degrees
 - cyan:    180 degrees
 - blue:    225 degrees
 - violet:  270 degrees
 - magenta: 315 degrees

The luminance is slightly lower than the luminance of fg2, to make fg2 pop out
a little bit.

For terminals, we map ANSI colors to editor colors as follows

 - black:   bg0 / fg0
 - red:     red / red
 - green:   green / green
 - yellow:  yellow / yellow
 - blue:    blue / blue
 - magenta: magenta / magenta
 - cyan:    cyan / cyan
 - white:   fg2 / fg4

Different colors have different perceived brightness (e.g., r:g:b is roughly
3:10:1) and terminals make use of that.  XTerm by default uses

 - red:     RGB(205,   0,   0), HCL( 12, 144, 43)
 - green:   RGB(  0, 205,   0), HCL(128, 111, 72)
 - yellow:  RGB(205, 205,   0), HCL( 86,  88, 80)
 - blue:    RGB(  0,   0, 205), HCL(266, 101, 25)
 - magenta: RGB(205,   0, 205), HCL(308, 111, 49)
 - cyan:    RGB(  0, 205, 205), HCL(192,  59, 75)

Red can indicate errors/deletions and is treated specially.  Blue is often
used for highlight and is treated specially as well.
'''

class Palette:
    GRAYSCALE_NAMES = [
        'bg0',
        'bg1',
        'bg2',
        'fg0',
        'fg1',
        'fg2',
        'fg3',
        'fg4',
    ]

    COLOR_NAMES = [
        'red',
        'orange',
        'yellow',
        'green',
        'cyan',
        'blue',
        'violet',
        'magenta',
    ]

    TERMINAL_OVERRIDES = [
        'red',
        'blue',
    ]

    def __init__(self, hue_offset, chroma, luminance):
        self.hue_offset = hue_offset
        self.chroma = chroma
        self.luminance = luminance

        self.grayscale = self.generate_grayscale()
        self.colors = self.generate_colors()

        self.palette = self.grayscale.copy()
        self.palette.update(self.colors)

    def name(self):
        return f'hcl{self.hue_offset:02}{self.chroma:02}{self.luminance:02}'

    def generate_grayscale(self):
        lumi_map = {
            'bg0': 0,
            'bg1': 10,
            'bg2': 30,
            'fg0': 50,
            'fg1': 70,
            'fg2': 80,
            'fg3': 90,
            'fg4': 100,
        }

        grayscale = {}
        for name in self.GRAYSCALE_NAMES:
            lumi = lumi_map[name]
            rgb = self.hcl_to_rgb(0, 0, lumi)
            code = self.rgb_to_hex(rgb)
            grayscale[name] = code

        return grayscale

    def generate_colors(self):
        hue_dist = 360 / len(self.COLOR_NAMES)
        colors = {}
        for i, name in enumerate(self.COLOR_NAMES + self.TERMINAL_OVERRIDES):
            if i >= len(self.COLOR_NAMES):
                i = self.COLOR_NAMES.index(name)
                name = 'terminal-' + name

            hue = self.hue_offset + hue_dist * i
            chroma = self.chroma
            lumi = self.luminance

            # overrides for terminals
            if name == 'terminal-red':
                # red is often used for errors/deletions
                hue += 12
                chroma += 50
                lumi -= 15
            elif name == 'terminal-blue':
                # blue is often used for highlight
                lumi -= 20

            chroma = self.fix_chroma(hue, chroma, lumi)
            print(f'{name}: HCL({hue}, {chroma}, {lumi})')

            rgb = self.hcl_to_rgb(hue, chroma, lumi)
            code = self.rgb_to_hex(rgb)
            colors[name] = code

        return colors

    def fix_chroma(self, hue, chroma, lumi):
        for chroma in range(chroma, -1, -1):
            rgb = self.hcl_to_rgb(hue, chroma, lumi)
            if rgb.max() <= 1.0 and rgb.min() >= 0.0:
                return chroma

    def hcl_to_rgb(self, hue, chroma, lumi):
        hcl = np.array([hue, chroma, lumi])
        lch = np.flip(hcl)
        luv = colour.LCHuv_to_Luv(lch)
        xyz = colour.Luv_to_XYZ(luv)
        rgb = colour.XYZ_to_sRGB(xyz)
        return rgb

    def rgb_to_hcl(self, rgb):
        xyz = colour.sRGB_to_XYZ(rgb)
        luv = colour.XYZ_to_Luv(xyz)
        lch = colour.Luv_to_LCHuv(luv)
        hcl = np.flip(lch)
        return hcl

    def rgb_to_hex(self, rgb):
        rgb = rgb.clip(0.0, 1.0)
        val = (rgb * 255.0).round().astype('B')
        code = val.tobytes().hex()
        return code

    def hex_to_rgb(self, code):
        val = np.frombuffer(bytes.fromhex(code), dtype='B')
        rgb = val.astype(float) / 255.0
        return rgb

def export_xresources(pal):
    ansi_map = [
        # normal
        'bg0',
        'terminal-red',
        'green',
        'yellow',
        'cyan',
        'terminal-blue',
        'magenta',
        'fg2',
        # bright
        'fg0',
        'terminal-red',
        'green',
        'yellow',
        'cyan',
        'terminal-blue',
        'magenta',
        'fg4',
    ]

    for i, name in enumerate(ansi_map):
        code = pal.palette[name]
        print(f'*.color{i}: #{code}')

    print(f'*.background: #{pal.palette[ansi_map[0]]}')
    print(f'*.foreground: #{pal.palette[ansi_map[7]]}')
    print()

def export_alacritty(pal):
    ansi_map = {
        'black':   'bg0',
        'red':     'terminal-red',
        'green':   'green',
        'yellow':  'yellow',
        'cyan':    'cyan',
        'blue':    'terminal-blue',
        'magenta': 'magenta',
        'white':   'fg2',
    }

    scheme = {
        'primary': {
        },
        'normal': {
        },
        'bright': {
        },
    }

    for ansi, name in ansi_map.items():
        code = '#' + pal.palette[name]
        scheme['normal'][ansi] = code
        scheme['bright'][ansi] = code

    # fix bright black/white
    scheme['bright']['black'] = '#' + pal.palette['fg0']
    scheme['bright']['white'] = '#' + pal.palette['fg4']

    # set background/foreground
    scheme['primary']['background'] = scheme['normal']['black']
    scheme['primary']['foreground'] = scheme['normal']['white']

    alacritty_yaml = {
        'color-schemes': {
            pal.name(): scheme
        },
    }
    print(dump(alacritty_yaml))

def export_vim(pal):
    # basead on how base16-shell maps ANSI colors
    base16_map = {
        'base00': 'bg0',     # ANSI black
        'base01': 'bg1',
        'base02': 'bg2',
        'base03': 'fg0',     # ANSI bright black
        'base04': 'fg1',
        'base05': 'fg2',     # ANSI white
        'base06': 'fg3',
        'base07': 'fg4',     # ANSI bright white

        'base08': 'red',     # ANSI normal and brigt red
        'base09': 'orange',
        'base0A': 'yellow',  # ANSI normal and brigt yellow
        'base0B': 'green',   # ANSI normal and brigt green
        'base0C': 'cyan',    # ANSI normal and brigt cyan
        'base0D': 'blue',    # ANSI normal and brigt blue
        'base0E': 'magenta', # ANSI normal and brigt magenta
        'base0F': 'violet',
    }

    for base16, name in base16_map.items():
        code = pal.palette[name]
        print(f'let s:gui{base16[4:]}        = "{code}"')

pal = Palette(0, 70, 70)
export_xresources(pal)
export_alacritty(pal)
export_vim(pal)
