# coding=utf-8
'''
python gen_cannon.py [img_n] [show_each]
    img_n: integer, generate img_n cannon images with random color
    show_each:
        1, show each image
        0, O/W

e.g.
$ python gen_cannon.py 3 1
generate 3 cannon images with random color and show each image.
'''
import sys
import os

import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt

W = 512
H = 512
margin = 50
cannon_w = 100  # 砲管寬
cannon_h = 300  # 砲管高
radius = 75  # 兩大圓圈
# 內部小圈
cir_r = 10  # 半徑
cir_n = radius//cir_r  # 數量
# 大砲前端
head_bottom = 150  # 寬
head_top = 120  # 寬
head_h = 125  # 上半部分高度
head_lh = 20  # 下半部分的高度


def rand_color():
    return np.random.randint(0, 256, size=3, dtype=np.uint8).tolist()


def shadow(color):
    return [round(i*0.4) for i in color]  # power-law


def gen_cannon(name, bg_color=[255, 255, 255], c_color=[255, 488, 242]):

    img = np.zeros((W, H, 3), np.uint8)  # BGR
    img[:] = bg_color
    color = c_color

    # 座標最左上是 (0,0)
    mid = W//2
    # 底座
    cv.rectangle(img, (margin, H-margin+1), (W-margin, H-1), (0, 0, 0), -1)
    # 砲管
    connon_pos_y = H-margin-cannon_h
    cv.rectangle(img, (mid-cannon_w//2, connon_pos_y),
                 (mid+cannon_w//2, H-margin), color, -1)  # 左上和右下的點
    # shadow
    pts2 = np.array([[mid, H-margin],
                     [mid+cannon_w//2, H-margin],
                     [mid+cannon_w//2, connon_pos_y]], np.int32)
    cv.fillPoly(img, pts2.reshape(-1, 3, 2), shadow(color))

    # 兩圓圈
    cv.circle(img, (mid-cannon_w//2-1-radius, H-margin-radius),
              radius, color, -1)
    cv.circle(img, (mid+cannon_w//2+1+radius, H-margin-radius),
              radius, color, -1)
    # 內部小圈
    for i, rc in zip(range(1, cir_n), [rand_color() for i in range(cir_n-1)]):
        cv.circle(img, (mid-cannon_w//2-1-radius, H-margin-radius),
                  radius-i*cir_r, rc, -1)
        cv.circle(img, (mid+cannon_w//2+1+radius, H-margin-radius),
                  radius-i*cir_r, rc, -1)

    # 大砲前端
    # 左上、右上、右下、左下
    pts1 = np.array([[mid-head_top//2, connon_pos_y-head_h-head_lh-1],
                     [mid+head_top//2, connon_pos_y-head_h-head_lh-1],
                     [mid+head_bottom//2, connon_pos_y-head_lh-1],
                     [mid-head_bottom//2, connon_pos_y-head_lh-1]], np.int32)
    pts2 = np.array([[mid-head_bottom//2, connon_pos_y-head_lh],
                     [mid+head_bottom//2, connon_pos_y-head_lh],
                     [mid+head_top//2, connon_pos_y-1],
                     [mid-head_top//2, connon_pos_y-1]], np.int32)  # shadow1
    pts3 = np.array([[mid+head_bottom//4, connon_pos_y-head_lh-1],
                     [mid+head_bottom//2, connon_pos_y-head_lh-1],
                     [mid+head_top//2, connon_pos_y-head_h-head_lh-1]],
                    np.int32)  # shadow2

    # cv.polylines(img, pts1.reshape((-1, 4, 2)), True, color, 5) # can't fill with color (thickness=-1)
    cv.fillPoly(img, pts1.reshape((-1, 4, 2)), color)
    cv.fillPoly(img, pts2.reshape((-1, 4, 2)), shadow(color))
    cv.fillPoly(img, pts3.reshape((-1, 3, 2)), shadow(color))

    return img


def show(img, name):
    # show
    p_img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    plt.figure(num="The Armstrong Cannon")
    # plt.xticks([]), plt.yticks([])
    plt.title(name)
    plt.tight_layout()
    plt.imshow(p_img)
    plt.show()


def main():
    n, show_each = 1, False
    if len(sys.argv) > 1:
        try:
            n = int(sys.argv[1])
        except:
            print("Not a number!! So I just give it one~")
        if len(sys.argv) > 2:
            try:
                show_each = int(sys.argv[2])
            except:
                print("I guess that you don't want to show each image.")

    for i in range(n):
        bg_color = rand_color()
        c_color = rand_color()
        name = "(Gintama)cannon"+str(i)+".jpg"
        img = gen_cannon(name, c_color=c_color)
        cv.imwrite(os.path.join(os.getcwd(), name), img)
        if show_each:
            show(img, name)

    print("-"*20, "generate "+str(n)+" image(s)", sep='\n')


if __name__ == "__main__":
    main()