#!/opt/python3.6/bin/python3.6
# -*- coding: utf-8 -*-

"""
PYNQ 上で pynqrouter を実行する。
"""

import argparse
import requests
import sys
import time
import threading

import BoardStr

from pynq import Overlay
from pynq import PL
from pynq import MMIO
from pynq.overlays.router04 import SolverOverlay

# Settings -- pynqrouter
OFFSET_BOARD_LOW   = 0x10000   # 0x10000 ~ 0x1ffff
OFFSET_BOARD_HIGH  = 0x20000   # 0x20000 ~ 0x2ffff
OFFSET_SEED   = 0x30000  # 0x30000
OFFSET_STATUS = 0x30008  # 0x30008
MAX_X = 72
MAX_Y = 72
MAX_Z = 8
BITWIDTH_Z = 3
BOARD_WORDS = (MAX_X * MAX_Y * MAX_Z) // 4

# Settings -- LED
# IP_LED = 'SEG_axi_gpio_0_Reg'

BS_IP = [
    ('/home/root/adc2018/PYNQ/pynq/overlays/router04.bit',  'pynqrouter_0')
]

solver_thread = None

class StoppableThread(threading.Thread):

    def __init__(self, target, args=()):
        super(StoppableThread, self).__init__(target=target, args=args)
        self._status = 'running'
    def stop(self):
        if self._status == 'running':
            self._status = 'stopping'
    def stopped(self):
        self._status = 'stopped'
    def is_running(self):
        return (self._status == 'running')
    def is_stopping(self):
        return (self._status == 'stopping')
    def is_stopped(self):
        return (self._status == 'stopped')

def start_solver(boardstr, seed, option):
    global solver_thread

    if solver_thread is None:
        solver_thread = StoppableThread(target=solve, args=(boardstr, int(seed), False, option))
        solver_thread.start()

def stop_solver():
    global solver_thread

    if solver_thread is not None:
        if solver_thread.is_running():
            solver_thread.stop()
            solver_thread.join()
            solver_thread = None

def solve(boardstr, seed=12345, zero_padding=False, option=dict()):
    global solver_thread

    print('boardstr:')
    print(boardstr)
    print('seed:')
    print(seed)
    print('')

    res = {
        'client': option['client'],
        'qname': option['name'],
        'problem': option['problem'],
        'solution': '',
        'cputime': -1.0,
        'req_id': option['req_id'],
        'solved': False
    }

    # LINE数を数えてコンフィグするbitstreamを分岐
    line_num = boardstr.count('L')
    if line_num < 1023:
        BITSTREAM = BS_IP[0][0]
        IP = BS_IP[0][1]
    else:
        solver_thread.stopped()
        solver_thread = None
        if "host" in option:
            r = requests.post("http://{}/post".format(option['host']), data=res)
        return res

    # ボード文字列から X, Y, Z を読んでくる
    size_x = (ord(boardstr[1]) - ord('0')) * 10 + (ord(boardstr[2]) - ord('0'))
    size_y = (ord(boardstr[4]) - ord('0')) * 10 + (ord(boardstr[5]) - ord('0'))
    size_z = (ord(boardstr[7]) - ord('0'))

    # Overlay 読み込み
    OL = SolverOverlay(BITSTREAM)
    print(OL.ip_dict)
    print('Overlay loaded!')

    # MMIO 接続 (pynqrouter)
    mmio = MMIO(OL.ip_dict[IP]['phys_addr'], OL.ip_dict[IP]['addr_range'])
    print("MMIO is configured")
    time.sleep(0.1)

    # MMIO 接続 & リセット (LED)
    # if isinstance(PL.ip_dict[IP_LED][0], int):
    #     mmio_led = MMIO(PL.ip_dict[IP_LED][0], PL.ip_dict[IP_LED][1])
    # else:
    #     mmio_led = MMIO(int(PL.ip_dict[IP_LED][0], 16), int(PL.ip_dict[IP_LED][1], 16))
    # mmio_led.write(0, 0)

    # 入力データをセット
    imem = pack(boardstr)
    for i in range(len(imem)):
        time.sleep(0.0001)
        mmio.write(OFFSET_BOARD_LOW + (i * 4), imem[i])
    print("Boardstr is sent")


    mmio.write(OFFSET_SEED, seed)

    # スタート
    # ap_start (0 番地の 1 ビット目 = 1)
    mmio.write(0, 1)
    print('Start!')
    time_start = time.time()

    # ap_done (0 番地の 2 ビット目 = 2) が立つまで待ってもいいが
    #   done は一瞬だけ立つだけのことがあるから
    # ap_idle (0 番地の 3 ビット目 = 4) を待ったほうが良い
    iteration = 0
    while (mmio.read(0) & 4) == 0:
        # 動いてるっぽく見えるようにLチカさせる
        iteration += 1
        if iteration == 10000:
            # mmio_led.write(0, 3)
            pass
        elif 20000 <= iteration:
            # mmio_led.write(0, 12)
            iteration = 0
        if (solver_thread is not None) and (not solver_thread.is_running()):
            solver_thread.stopped()
            solver_thread = None
            if "host" in option:
                r = requests.post("http://{}/post".format(option['host']), data=res)
            return res

    # 完了の確認
    print('Done!')
    print('control:', mmio.read(0))
    time_done = time.time()
    elapsed = time_done - time_start
    print('elapsed:', elapsed)
    print('')

    # 状態の取得
    status = int(mmio.read(OFFSET_STATUS))
    print('status:', status)
    if status != 0:
        # 解けなかったらLEDを消す
        # mmio_led.write(0, 0)
        sys.stderr.write('Cannot solve it!\n')
        if "host" in option:
            r = requests.post("http://{}/post".format(option['host']), data=res)
        solver_thread = None
        return res
    
    print('Solved!')

    # 解けたらLEDを全部つける
    # mmio_led.write(0, 15)

    # 出力
    omem = []
    for i in range(BOARD_WORDS):
        omem.append(mmio.read(OFFSET_BOARD_LOW + (i * 4)))
    boards_low = unpack(omem)

    omem = []
    for i in range(BOARD_WORDS):
        omem.append(mmio.read(OFFSET_BOARD_HIGH + (i * 4)))
    boards_high = unpack(omem)

    boards = []
    for _high, _low in zip(boards_high, boards_low):
        boards.append(_high * 256 + _low)

    # 回答の生成
    solution = ('SIZE ' + str(size_x) + 'X' + str(size_y) + 'X' + str(size_z) + '\n')
    for z in range(size_z):
        solution += ('LAYER ' + str(z + 1) + '\n')
        for y in range(size_y):
            for x in range(size_x):
                if x != 0:
                    solution += ','
                i = ((x * MAX_X + y) << BITWIDTH_Z) | z
                if zero_padding:
                    solution += '{0:0>2}'.format(boards[i])  # 2桁の0詰め
                else:
                    solution += str(boards[i])               # 普通に表示
            solution += '\n'

    res = {
        'client': option['client'],
        'qname': option['name'],
        'problem': option['problem'],
        'solution': solution,
        'cputime': elapsed,
        'req_id': option['req_id'],
        'solved': True
    }
    if "resolver" in option:
        r = requests.post("http://{}/post".format(option['resolver']), data=res)
    elif "host" in option:
        r = requests.post("http://{}/post".format(option['host']), data=res)
 
    if solver_thread is not None:
        solver_thread.stopped()
        solver_thread = None

    return res

def main():
    parser = argparse.ArgumentParser(description="Solver with pynqrouter")

    # 入出力
    parser.add_argument('-i', '--input', action='store', nargs='?', default=None, type=str,
        help='Path to input problem file')
    parser.add_argument('-b', '--boardstr', action='store', nargs='?', default=None, type=str,
        help='Problem boardstr (if you want to solve directly)')
    parser.add_argument('-o', '--output', action='store', nargs='?', default=None, type=str,
        help='Path to output answer file')

    # ソルバオプション
    parser.add_argument('-t', '--terminals', action='store', nargs='?', default='initial', type=str,
        help='Terminals order (initial, edgefirst, or random) (default: initial)')
    parser.add_argument('-s', '--seed', action='store', nargs='?', default=12345, type=int,
        help='Random seed')

    # 出力オプション
    parser.add_argument('-p', '--print', action='store_true', default=False,
        help='Enable to print the solution')
    parser.add_argument('-z', '--zero-padding', action='store_true', default=False,
        help='Enable to do zero-padding')

    args = parser.parse_args()

    if args.input is not None:
        # 問題ファイルの読み込み
        with open(args.input, 'r') as f:
            lines = f.readlines()

        # 問題ファイルを boardstr に変換
        boardstr = BoardStr.conv_boardstr(lines, args.terminals, args.seed)

    elif args.boardstr is not None:
        boardstr = args.boardstr

    else:
        sys.stderr.write('Specify at least "input" or "boardstr"!\n')
        sys.exit(1)

    opt = {
        'client': "localhost",
        'host': "localhost",
        'name': "",
        'problem': "",
        'req_id': "",
    }

    result = solve(boardstr, args.seed, option=opt)
    if result['solved'] == False:
        sys.exit(1)

    # 表示 & ファイル出力
    if args.print:
        print('')
        print('SOLUTION')
        print('========')
        print(result['solution'])
    if args.output is not None:
        with open(args.output, 'w') as f:
            f.write(result['solution'])


def pack(_str):
    """
    ボードストリング文字列をMMIOメモリにパックする
    """
    _mem = [ 0 for _ in range(MAX_X * MAX_Y * MAX_Z // 4) ]
    for i in range(len(_str)):
        _index  = i // 4;
        _offset = i % 4;
        if _offset == 0:
            _mem[_index] = _mem[_index] | (ord(_str[i]))
        elif _offset == 1:
            _mem[_index] = _mem[_index] | (ord(_str[i]) << 8)
        elif _offset == 2:
            _mem[_index] = _mem[_index] | (ord(_str[i]) << 16)
        elif _offset == 3:
            _mem[_index] = _mem[_index] | (ord(_str[i]) << 24)
    return _mem


def unpack(_mem):
    """
    MMIOメモリからボード情報をアンパックする
    """
    _boards = [ 0 for _ in range(MAX_X * MAX_Y * MAX_Z) ]
    for i in range(len(_mem)):
        _boards[4 * i    ] = (_mem[i])       & 255
        _boards[4 * i + 1] = (_mem[i] >> 8)  & 255
        _boards[4 * i + 2] = (_mem[i] >> 16) & 255
        _boards[4 * i + 3] = (_mem[i] >> 24) & 255
    return _boards


if __name__ == '__main__':
    main()