import argparse
import yaml
import math
import os
import struct
import torch
from utils.torch_utils import select_device


class YoloLayers():
    def get_route(self, n, layers):
        route = 0
        for i, layer in enumerate(layers):
            if i <= n:
                route += layer[1]
            else:
                break
        return route

    def route(self, layers=''):
        return '\n[route]\n' + \
               'layers=%s\n' % layers

    def reorg(self):
        return '\n[reorg]\n'

    def shortcut(self, route=-1, activation='linear'):
        return '\n[shortcut]\n' + \
               'from=%d\n' % route + \
               'activation=%s\n' % activation

    def maxpool(self, stride=1, size=1):
        return '\n[maxpool]\n' + \
               'stride=%d\n' % stride + \
               'size=%d\n' % size

    def upsample(self, stride=1):
        return '\n[upsample]\n' + \
               'stride=%d\n' % stride

    def convolutional(self, bn=False, size=1, stride=1, pad=1, filters=1, groups=1, activation='linear'):
        b = 'batch_normalize=1\n' if bn is True else ''
        g = 'groups=%d\n' % groups if groups > 1 else ''
        return '\n[convolutional]\n' + \
               b + \
               'filters=%d\n' % filters + \
               'size=%d\n' % size + \
               'stride=%d\n' % stride + \
               'pad=%d\n' % pad + \
               g + \
               'activation=%s\n' % activation

    def yolo(self, mask='', anchors='', classes=80, num=3):
        return '\n[yolo]\n' + \
               'mask=%s\n' % mask + \
               'anchors=%s\n' % anchors + \
               'classes=%d\n' % classes + \
               'num=%d\n' % num + \
               'scale_x_y=2.0\n' + \
               'beta_nms=0.6\n' + \
               'new_coords=1\n'


def parse_args():
    parser = argparse.ArgumentParser(description='PyTorch YOLOv5 conversion')
    parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
    parser.add_argument('-c', '--yaml', help='Input cfg (.yaml) file path')
    parser.add_argument(
        '-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
    args = parser.parse_args()
    if not os.path.isfile(args.weights):
        raise SystemExit('Invalid weights file')
    if not args.yaml:
        args.yaml = ''
    return args.weights, args.yaml, args.size


def get_width(x, gw, divisor=8):
    return int(math.ceil((x * gw) / divisor)) * divisor


def get_depth(x, gd):
    if x == 1:
        return 1
    r = int(round(x * gd))
    if x * gd - int(x * gd) == 0.5 and int(x * gd) % 2 == 0:
        r -= 1
    return max(r, 1)


pt_file, yaml_file, inference_size = parse_args()

model_name = os.path.basename(pt_file).split('.pt')[0]
wts_file = model_name + '.wts' if 'yolov5' in model_name else 'yolov5_' + model_name + '.wts'
cfg_file = model_name + '.cfg' if 'yolov5' in model_name else 'yolov5_' + model_name + '.cfg'

if yaml_file == '':
    yaml_file = 'models/' + model_name + '.yaml'
    if not os.path.isfile(yaml_file):
        yaml_file = 'models/hub/' + model_name + '.yaml'
        if not os.path.isfile(yaml_file):
            raise SystemExit('YAML file not found')
elif not os.path.isfile(yaml_file):
    raise SystemExit('Invalid YAML file')

device = select_device('cpu')
model = torch.load(pt_file, map_location=device)['model'].float()

anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]
delattr(model.model[-1], 'anchor_grid')
model.model[-1].register_buffer('anchor_grid', anchor_grid)

model.to(device).eval()

nc = 0
anchors = ''
masks = []

yolo_idx = 0
spp_idx = 0

for k, v in model.state_dict().items():
    if 'anchor_grid' in k:
        yolo_idx = int(k.split('.')[1])
        vr = v.cpu().numpy().tolist()
        a = v.reshape(-1).cpu().numpy().astype(float).tolist()
        anchors = str(a)[1:-1]
        num = 0
        for m in vr:
            mask = []
            for _ in range(len(m)):
                mask.append(num)
                num += 1
            masks.append(mask)
    elif '.%d.m.0.weight' % yolo_idx in k:
        vr = v.cpu().numpy().tolist()
        nc = int((len(vr) / len(masks[0])) - 5)

with open(cfg_file, 'w') as c:
    with open(yaml_file, 'r', encoding='utf-8') as f:
        c.write('[net]\n')
        c.write('width=%d\n' % (inference_size[0] if len(inference_size) == 1 else inference_size[1]))
        c.write('height=%d\n' % inference_size[0])
        c.write('channels=3\n')
        c.write('letter_box=1\n')
        depth_multiple = 0
        width_multiple = 0
        layers = []
        yoloLayers = YoloLayers()
        f = yaml.load(f, Loader=yaml.FullLoader)
        for topic in f:
            if topic == 'depth_multiple':
                depth_multiple = f[topic]
            elif topic == 'width_multiple':
                width_multiple = f[topic]
            elif topic == 'backbone' or topic == 'head':
                for v in f[topic]:
                    if v[2] == 'Focus':
                        layer = '\n# Focus\n'
                        blocks = 0
                        layer += yoloLayers.reorg()
                        blocks += 1
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1],
                                                          activation='silu')
                        blocks += 1
                        layers.append([layer, blocks])
                    if v[2] == 'Conv':
                        layer = '\n# Conv\n'
                        blocks = 0
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1],
                                                          stride=v[3][2], activation='silu')
                        blocks += 1
                        layers.append([layer, blocks])
                    elif v[2] == 'C3':
                        layer = '\n# C3\n'
                        blocks = 0
                        # SPLIT
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                          activation='silu')
                        blocks += 1
                        layer += yoloLayers.route(layers='-2')
                        blocks += 1
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                          activation='silu')
                        blocks += 1
                        # Residual Block
                        if len(v[3]) == 1 or v[3][1] is True:
                            for _ in range(get_depth(v[1], depth_multiple)):
                                layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                                  activation='silu')
                                blocks += 1
                                layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                                  size=3, activation='silu')
                                blocks += 1
                                layer += yoloLayers.shortcut(route=-3)
                                blocks += 1
                            # Merge
                            layer += yoloLayers.route(layers='-1, -%d' % (3 * get_depth(v[1], depth_multiple) + 3))
                            blocks += 1
                        else:
                            for _ in range(get_depth(v[1], depth_multiple)):
                                layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                                  activation='silu')
                                blocks += 1
                                layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                                  size=3, activation='silu')
                                blocks += 1
                            # Merge
                            layer += yoloLayers.route(layers='-1, -%d' % (2 * get_depth(v[1], depth_multiple) + 3))
                            blocks += 1
                        # Transition
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
                                                          activation='silu')
                        blocks += 1
                        layers.append([layer, blocks])
                    elif v[2] == 'SPP':
                        spp_idx = len(layers)
                        layer = '\n# SPP\n'
                        blocks = 0
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                          activation='silu')
                        blocks += 1
                        layer += yoloLayers.maxpool(size=v[3][1][0])
                        blocks += 1
                        layer += yoloLayers.route(layers='-2')
                        blocks += 1
                        layer += yoloLayers.maxpool(size=v[3][1][1])
                        blocks += 1
                        layer += yoloLayers.route(layers='-4')
                        blocks += 1
                        layer += yoloLayers.maxpool(size=v[3][1][2])
                        blocks += 1
                        layer += yoloLayers.route(layers='-6, -5, -3, -1')
                        blocks += 1
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
                                                          activation='silu')
                        blocks += 1
                        layers.append([layer, blocks])
                    elif v[2] == 'SPPF':
                        spp_idx = len(layers)
                        layer = '\n# SPPF\n'
                        blocks = 0
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
                                                          activation='silu')
                        blocks += 1
                        layer += yoloLayers.maxpool(size=v[3][1])
                        blocks += 1
                        layer += yoloLayers.maxpool(size=v[3][1])
                        blocks += 1
                        layer += yoloLayers.maxpool(size=v[3][1])
                        blocks += 1
                        layer += yoloLayers.route(layers='-4, -3, -2, -1')
                        blocks += 1
                        layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
                                                          activation='silu')
                        blocks += 1
                        layers.append([layer, blocks])
                    elif v[2] == 'nn.Upsample':
                        layer = '\n# nn.Upsample\n'
                        blocks = 0
                        layer += yoloLayers.upsample(stride=v[3][1])
                        blocks += 1
                        layers.append([layer, blocks])
                    elif v[2] == 'Concat':
                        route = v[0][1]
                        route = yoloLayers.get_route(route, layers) if route > 0 else \
                            yoloLayers.get_route(len(layers) + route, layers)
                        layer = '\n# Concat\n'
                        blocks = 0
                        layer += yoloLayers.route(layers='-1, %d' % (route - 1))
                        blocks += 1
                        layers.append([layer, blocks])
                    elif v[2] == 'Detect':
                        for i, n in enumerate(v[0]):
                            route = yoloLayers.get_route(n, layers)
                            layer = '\n# Detect\n'
                            blocks = 0
                            layer += yoloLayers.route(layers='%d' % (route - 1))
                            blocks += 1
                            layer += yoloLayers.convolutional(filters=((nc + 5) * len(masks[i])), activation='logistic')
                            blocks += 1
                            layer += yoloLayers.yolo(mask=str(masks[i])[1:-1], anchors=anchors, classes=nc, num=num)
                            blocks += 1
                            layers.append([layer, blocks])
        for layer in layers:
            c.write(layer[0])

with open(wts_file, 'w') as f:
    wts_write = ''
    conv_count = 0
    cv1 = ''
    cv3 = ''
    cv3_idx = 0
    for k, v in model.state_dict().items():
        if 'num_batches_tracked' not in k and 'anchors' not in k and 'anchor_grid' not in k:
            vr = v.reshape(-1).cpu().numpy()
            idx = int(k.split('.')[1])
            if '.cv1.' in k and '.m.' not in k and idx != spp_idx:
                cv1 += '{} {} '.format(k, len(vr))
                for vv in vr:
                    cv1 += ' '
                    cv1 += struct.pack('>f', float(vv)).hex()
                cv1 += '\n'
                conv_count += 1
            elif cv1 != '' and '.m.' in k:
                wts_write += cv1
                cv1 = ''
            if '.cv3.' in k:
                cv3 += '{} {} '.format(k, len(vr))
                for vv in vr:
                    cv3 += ' '
                    cv3 += struct.pack('>f', float(vv)).hex()
                cv3 += '\n'
                cv3_idx = idx
                conv_count += 1
            elif cv3 != '' and cv3_idx != idx:
                wts_write += cv3
                cv3 = ''
                cv3_idx = 0
            if '.cv3.' not in k and not ('.cv1.' in k and '.m.' not in k and idx != spp_idx):
                wts_write += '{} {} '.format(k, len(vr))
                for vv in vr:
                    wts_write += ' '
                    wts_write += struct.pack('>f', float(vv)).hex()
                wts_write += '\n'
                conv_count += 1
    f.write('{}\n'.format(conv_count))
    f.write(wts_write)