from glob import glob
from shutil import copyfile
from json import load
from random import shuffle
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# stat
stat = {"file": []}
all_shapes = []
str_for_map = "0123456789.:-"
label_map = {j: i for (i, j) in enumerate(str_for_map)}
dst = "train"
src = "ori"
train_path = "/home/tmp/"


# plot
def plot_imgandlabel(img, txt):
    plt.imshow(img)
    cord = []

    for labels in txt:
        h, w = img.shape[:2]
        x1 = w * (labels[1] - labels[3]/2)
        y1 = h * (labels[2] - labels[4]/2)
        x2 = w * (labels[1] + labels[3]/2)
        y2 = h * (labels[2] + labels[4]/2)
        print(labels[0], x1, y1, x2, y2)

        bbox = patches.Rectangle((x1, y1), x2-x1, y2-y1,
            linewidth=2, facecolor='none', edgecolor='blue')
        plt.gca().add_patch(bbox)

    plt.imshow(img)
    plt.show()


# plot by name
def plot_imgandlabel_raw(imgpath, txtpath):
    img = plt.imread(imgpath)
    txt = open(txtpath).readlines()
    labels = [[float(t) for t in ti.split()] for ti in txt]
    plot_imgandlabel(img, labels)


# copy imgae to train
def copyImage(base):
    imgfile = base + ".jpg"
    copyfile(imgfile, imgfile.replace(src, dst))


# transfer label to train
def copyLabel(base):
    stat["file"].append(base)
    return
    js = load(open(base + ".json"))

    # image meta
    h, w = js["imageHeight"], js["imageWidth"]
    # debug
    # print(h, w)
    # imgfile = base + ".jpg"
    # img = plt.imread(imgfile)

    # shape transfer
    txt = []
    for sh in js["shapes"]:
        label, (x1, y1), (x2, y2)= sh['label'], *sh["points"]
        x1, x2 = min(x1, x2), max(x1, x2)
        y1, y2 = min(y1, y2), max(y1, y2)
        all_shapes.append((x2 - x1, y2 - y1))
        labels = [label_map[label], (x1 + x2) / 2 / w, (y1 + y2) / 2 / h, (x2 - x1) / w, (y2 - y1) / h]
        txt.append(labels)
        # stat
        if not stat.get(label):
            stat[label] = 1
        else:
            stat[label] += 1

    # write
    fout = open(base.replace(src, dst) + ".txt", "w")
    for t in txt:
        print("{} {:6f} {:6f} {:6f} {:6f}".format(*t), file=fout)
    # debug
    # print(txt)
    # plot_imgandlabel(img, txt)
    # break


def writeMeta():
    data = """\
classes={}
train=../train/train.txt
valid=../train/valid.txt
names=../train/name.txt
""".format(len(str_for_map))
    open(dst + "/name.txt", "w").write("\n".join(str_for_map) + "\n")
    open(dst + "/digit.data", "w").write(data)
    all_file = [train_path + d.replace(src, dst) + ".jpg" for d in stat['file']]
    # random
    shuffle(all_file)
    vlen = len(all_file) // 5
    train_file = all_file[:-vlen]
    valid_file = all_file[-vlen:]
    open(dst + "/train.txt", "w").write(
            "\n".join(train_file))
    open(dst + "/valid.txt", "w").write(
            "\n".join(valid_file))