#!/usr/bin/env python3
# Copyright 2020 vg
# SPDX-License-Identifier: MIT

'''
Usage: solver.py -h|--help
       solver.py [-p] YAML_GRID
       solver.py --print YAML_GRID


Options:
  -h, --help     Display this help message
  -p             Display solution with parents [DEFAULT is without parents]
  --print        Display only the original (unresolved) grid
'''

import copy
import os
import pickle
import time

import docopt
import yaml

class Kana:

    types = ('void', 'norm', 'froz', 'rock', 'myst', 'slim',
             'ar_u', 'ar_r', 'ar_d', 'ar_l')
    slime_merge_types = ('norm', 'froz', 'ar_u', 'ar_r', 'ar_d', 'ar_l')

    cache = {}

    def __init__(self, type_name, kana=None):
        self.type_name = type_name
        self.kana = kana

        #print(type_name)
        #print(kana)
        assert type_name in self.types

        # missing yet: 'n' and 'y...' and 'w...' are not we all vowels
        if kana:
            if type_name == 'slim':
                assert kana[0] in 'aiueo'
            else:
                assert kana[0] in 'kstnhmryw'
                assert kana[1] in 'aiueo'

    def create(type_name, kana=None):
        key = (type_name, kana)
        if key in Kana.cache:
            return Kana.cache[key]
        new_kana = Kana.cache[key] = Kana(type_name, kana)
        return new_kana

    def dump(self):
        'return a list representing the members of this Kana'
        return [self.type_name, self.kana]

    def __repr__(self):
        return "%s(%s)" % (self.type_name, self.kana)

    def __eq__(self, other):
        return self.type_name == other.type_name and self.kana == other.kana

    def __hash__(self):
        #print(f'hash of Kana({self}): {hash((self.type_name, self.kana))}')
        return hash((self.type_name, self.kana))


kana_void = Kana.create('void')


class KanaGrid:

    actions = ('reveal', 'up', 'right', 'down', 'left')

    def __init__(self, size, grid,
                 action_count=0, score=0, myst_count=0, parent=None):

        self.width = size[0]
        self.height = size[1]
        self.grid = copy.copy(grid)
        self.action_count = action_count
        self.score = score
        self.myst_count = myst_count
        self.parent = parent

    def copy(self):
        new_grid = KanaGrid(
                (self.width, self.height),
                grid=self.grid,
                action_count=self.action_count,
                score=self.score,
                myst_count=self.myst_count,
                parent=self.parent,
        )
        return new_grid

    def is_swappable(self, pos1, pos2):
        kana1 = self.get_kana(pos1)
        kana2 = self.get_kana(pos2)
        table_ok = {
            'norm': ('norm', 'froz', 'ar_u', 'ar_r', 'ar_d', 'ar_l'),
            'froz': ('norm',         'ar_u', 'ar_r', 'ar_d', 'ar_l'),
            'ar_u': ('norm', 'froz',                 'ar_d'        ),
            'ar_r': ('norm', 'froz',                         'ar_l'),
            'ar_d': ('norm', 'froz', 'ar_u'                        ),
            'ar_l': ('norm', 'froz',         'ar_r'                ),
        }
        if kana1.type_name in table_ok:
            if kana2.type_name in table_ok[kana1.type_name]:

                # early return for empty norm to empty norm
                if kana1.type_name == kana2.type_name == 'norm' \
                        and kana1.kana is kana2.kana is None:
                    return False

                ar_vect_ok = {
                        'ar_u': ( 0, -1),
                        'ar_l': (-1,  0),
                        'ar_d': ( 0,  1),
                        'ar_r': ( 1,  0),
                }
                vect1_target = ar_vect_ok.get(kana1.type_name, None)
                vect2_target = ar_vect_ok.get(kana2.type_name, None)
                if vect1_target or vect2_target:
                    vect1 = (pos2[0] - pos1[0], pos2[1] - pos1[1])
                    vect2 = (pos1[0] - pos2[0], pos1[1] - pos2[1])
                    if vect1 != vect1_target and vect2 != vect2_target:
                        return False
                return True
        return False

    def action(self, pos, action_type):
        kana = self.get_kana(pos)
        if action_type == 'reveal':
            if kana.type_name == 'myst':
                new_grid = self.copy()
                new_grid.action_count += 1
                new_grid.set_kana(pos, Kana.create('norm', kana.kana))
                return new_grid
        elif action_type in ('up', 'right', 'down', 'left'):
            pos_map = {
                    'up':       (pos[0], pos[1]-1),
                    'right':    (pos[0]+1, pos[1]),
                    'down':     (pos[0], pos[1]+1),
                    'left':     (pos[0]-1, pos[1]),
            }
            pos_dest = pos_map[action_type]
            if kana.type_name == 'slim':
                kana2 = self.get_kana(pos_dest)
                if kana2.type_name in Kana.slime_merge_types and kana2.kana:
                    # early return for slime merging for same vowel
                    if kana2.kana[1] == kana.kana[0]:
                        return
                    new_grid = self.copy()
                    new_grid.action_count += 1
                    new_kana1 = kana_void
                    new_kana2 = Kana.create(
                            kana2.type_name,
                            kana2.kana[0] + kana.kana[0],
                    )
                    new_grid.set_kana(pos, new_kana1)
                    new_grid.set_kana(pos_dest, new_kana2)
                    return new_grid
            elif self.is_swappable(pos, pos_dest):
                new_grid = self.copy()
                new_grid.action_count += 1
                new_grid.swap_kana(pos, pos_dest)
                return new_grid

    def generate_valid_pos_for_chain(self):
        for y in range(self.height):
            for x in range(self.width):
                kana = self.get_kana((x, y))
                if kana.kana and kana.type_name != 'slim':
                    yield (x, y)

    def populate_chain(self, pos1, chain_positions):
        myst_count = 0
        if pos1 in chain_positions:
            return myst_count
        kana1 = self.get_kana(pos1)
        if kana1.type_name == 'myst':
            myst_count += 1
        chain_positions.add(pos1)
        pos2_list = [
                (pos1[0], pos1[1]-1), # up
                (pos1[0]+1, pos1[1]), # right
                (pos1[0], pos1[1]+1), # down
                (pos1[0]-1, pos1[1]), # left
        ]
        for pos2 in pos2_list:
            if pos2 in chain_positions:
                continue
            kana2 = self.get_kana(pos2)
            if kana2.kana and kana2.type_name != 'slim':
                if is_kana_compatible(kana1, kana2):
                    myst_count += self.populate_chain(pos2, chain_positions)
        return myst_count

    def longest_chain(self):
        already_evaluated_pos = set()
        highest_length = 0
        highest_length_chain = 0
        highest_myst_count = 0
        for pos in self.generate_valid_pos_for_chain():
            if pos in already_evaluated_pos:
                continue
            chain = set()
            myst_count = self.populate_chain(pos, chain)
            already_evaluated_pos = already_evaluated_pos.union(chain)
            if myst_count > highest_myst_count:
                highest_myst_count = myst_count
            if highest_length < len(chain):
                highest_length = len(chain)
                highest_length_chain = chain
        return highest_length, highest_myst_count
        #, highest_length_chain # easy to add if needed

    def update_score(self):
        if not self.score:
            # if self.score == 0 calculate the score as the score is always at
            # the bare minimum equal to 1.
            self.score, self.myst_count = self.longest_chain()

    def get_grid_hashable(self):
        return tuple(self.grid)

    def get_kana(self, pos):
        if pos[0] < 0 or pos[0] >= self.width:
            return kana_void
        elif pos[1] < 0 or pos[1] >= self.height:
            return kana_void

        return self.grid[pos[0]+pos[1]*self.width]

    def set_kana(self, pos, kana):
        if pos[0] < 0 or pos[0] >= self.width:
            return
        elif pos[1] < 0 or pos[1] >= self.height:
            return
        self.grid[pos[0]+pos[1]*self.width] = kana
        self.score = 0

    def swap_kana(self, pos1, pos2):
        kana_dst = self.get_kana(pos2)

        if kana_dst.type_name in ('froz', 'fblk'):
            pos_tmp = pos1
            pos1 = pos2
            pos2 = pos_tmp
            kana_dst = self.get_kana(pos2)

        kana_src = self.get_kana(pos1)
        pos_src = pos1
        pos_dst = pos2
        vect = (pos2[0] - pos1[0], pos2[1] - pos1[1])

        while self.is_swappable(pos_src, pos_dst):
            self.set_kana(pos_src, kana_dst)
            self.set_kana(pos_dst, kana_src)

            if kana_src.type_name not in ('froz', 'fblk') :
                break

            pos_src = pos_dst
            pos_dst = (pos_dst[0] + vect[0], pos_dst[1] + vect[1])
            kana_dst = self.get_kana(pos_dst)

    def load(input_dict):
        grid = [
            Kana.create(serialized_kana[0], serialized_kana[1])
            for serialized_kana in input_dict['grid']
        ]
        params = {
            'size': input_dict['size'],
            'grid': grid,
            'action_count': input_dict.get('action_count', 0),
            'score': input_dict.get('score', 0),
            'myst_count': input_dict.get('myst_count', 0),
        }
        return KanaGrid(**params)

    def load_reversed_parents(input_dict_list):
        last_kanagrid = None
        for input_dict in input_dict_list:
            kanagrid = KanaGrid.load(input_dict)
            kanagrid.parent = last_kanagrid
            last_kanagrid = kanagrid
        return last_kanagrid


    def dump(self):
        'returns a dict representation of the kanagrid'
        output_dict = {
            'size': [self.width, self.height],
            'action_count': self.action_count,
            'score': self.score,
            'myst_count': self.myst_count,
            'grid': [kana.dump() for kana in self.grid],
        }
        return output_dict

    def dump_reversed_parents(self):
        kana_grids = []
        grid = self
        while grid:
            kana_grids.append(grid.dump())
            grid = grid.parent
        return reversed(kana_grids)

    def __repr__(self):
        self.update_score()
        return (
                ('KanaGrid (cnt: %d, score: %d): \n  ' % (self.action_count, self.score))
                + '\n  '.join(repr_grid(self.grid, (self.width, self.height)).splitlines())
        )

    def __eq__(self, other):
        return (
                self.width == other.width
                and self.height == other.height
                and self.grid == other.grid
                and self.action_count == other.action_count
                )


def repr_grid(grid, grid_size):
    indicator_map = {
            'norm': (' ', ' '),
            'froz': ('\x1b[36m[', ']\x1b[0m'),
            'rock': (' \x1b[1;40m', '\x1b[0m '),
            'myst': ('\x1b[33m?', '?\x1b[0m'),
            'slim': ('\x1b[32m~', '~\x1b[0m'),
            'ar_u': ('\x1b[31m∧\x1b[0m', '\x1b[31m∧\x1b[0m'),
            'ar_r': ('\x1b[31m>\x1b[0m', '\x1b[31m>\x1b[0m'),
            'ar_d': ('\x1b[31m∨\x1b[0m', '\x1b[31m∨\x1b[0m'),
            'ar_l': ('\x1b[31m<\x1b[0m', '\x1b[31m<\x1b[0m'),
    }
    lines = []
    kana_iter = iter(grid)
    for y in range(grid_size[1]):
        line = ''
        for x in range(grid_size[0]):
            kana = next(kana_iter)
            tname = kana.type_name
            kkana = kana.kana
            if not kkana:
                kkana = '  '
            if tname == 'void':
                line += '      '
            elif tname in indicator_map:
                line += '|%s%2s%s|' % (
                        indicator_map[tname][0],
                        kkana,
                        indicator_map[tname][1],
                )
        lines.append(line)
    return '\n'.join(lines)


def display_grid(grid, grid_size):
    print(repr_grid(grid, grid_size))


def is_kana_compatible(kana1, kana2):
    if kana1.type_name == 'slim' or kana2.type_name == 'slim':
        return False
    if kana1.kana[0] == kana2.kana[0] or kana1.kana[1] == kana2.kana[1]:
        return True
    return False


def generate_possible_grids(kanagrid):
    for y in range(kanagrid.height):
        for x in range(kanagrid.width):
            for action_type in KanaGrid.actions:
                new_grid = kanagrid.action((x, y), action_type)
                if new_grid:
                    # better perf to have only new_grid is None check and not
                    # comparison if move has given same grid as having exact
                    # same grid is less likely than having action() filter
                    # useless actions and comparing all grid two times (one
                    # here and one in the taboos) is more consuming.
                    yield (x, y), action_type, new_grid


def generate_all_possible_grids(grid, bests, max_actions):
    for pos, action_type, new_grid in generate_possible_grids(grid):
        key = new_grid.get_grid_hashable()
        if key in bests and bests[key].action_count <= new_grid.action_count:
            continue
        bests[key] = new_grid
        new_grid.parent = grid
        if new_grid.action_count >= max_actions:
            yield new_grid
            continue
        yield from generate_all_possible_grids(new_grid, bests, max_actions)


def search_all_solution(kanagrid, target_score, max_actions):
    bests = {}
    generator = generate_all_possible_grids(kanagrid, bests=bests, max_actions=max_actions)
    for grid in generator:
        grid.update_score()
        if grid.score >= target_score and grid.myst_count == 0:
            yield grid


def get_solutions_from_cache(cache):
    for solution in cache['solutions']:
        kanagrid = KanaGrid.load_reversed_parents(solution)
        yield kanagrid


def repr_grid_with_parents(grid):
    items = []
    while grid:
        items.append(str(grid))
        grid = grid.parent
    return '\n'.join(reversed(items))


def main():

    args = docopt.docopt(__doc__)
    grid_fn = args['YAML_GRID']
    with open(grid_fn, encoding='utf8') as stream:
        input_dict = yaml.safe_load(stream)
    kanagrid = KanaGrid.load(input_dict)
    target_score = input_dict['target_score']
    max_actions = input_dict['max_actions']
    print('Size: %dx%d' % (kanagrid.width, kanagrid.height))
    print('Target score: %d' % target_score)
    print('Max actions: %d' % max_actions)
    print('Initial grid:')
    print(kanagrid)

    if args['--print']:
        return

    solver_start_time = int(time.time())
    mtime = os.path.getmtime
    cache_file = '%s.cache' % grid_fn

    if os.path.exists(cache_file) and mtime(cache_file) >= mtime(grid_fn):
        print(f'Using cached version from {cache_file}')
        with open(cache_file, 'rb') as stream:
            cached_version = pickle.load(stream)
        generator = get_solutions_from_cache(cached_version)
        cache_generation = None
    else:
        generator = search_all_solution(kanagrid, target_score, max_actions)
        cache_generation = {
            'max_actions': 1,
            'target_score': 1,
            'solutions': [],
        }

    for grid in generator:
        if cache_generation:
            cache_generation['solutions'].append(grid.dump_reversed_parents())
        print("="*80)
        if args['-p']:
            print(repr_grid_with_parents(grid))
        else:
            print(grid)

    if cache_generation:
        with open(cache_file, 'wb') as stream:
            pickle.dump(cache_generation, stream)

    solver_stop_time = int(time.time())
    solver_delta_time = solver_stop_time - solver_start_time
    hours = solver_delta_time // 3600
    minutes = (solver_delta_time // 60) % 60
    seconds = solver_delta_time % 60

    print(f'time taken to calculate: {hours:02d}:{minutes:02d}:{seconds:02d}')


if __name__ == '__main__':
    main()