# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
import time
from PIL import Image, ImageDraw, ImageFont
from datetime import datetime, timedelta
import shutil
pic_source_dir = os.path.join(os.getcwd(), "news_pic_source\\")  # 原图路径
pic_crop_dir = os.path.join(os.getcwd(), "news_pic_crop\\")  # 裁剪后的图片路径
pic_font_dir = os.path.join(os.getcwd(), "news_pic_font\\")  # 加字后的图片路径
start_date = "20190112"  # 绘制图片的其起始日期
# 判断文件夹是否存在，不存在则新建
def is_dir_existed(path, mkdir=True):
    if mkdir:
        if not os.path.exists(path):
            os.makedirs(path)
    else:
        return os.path.exists(path)
# opencv读取中文路径名会乱码
def cv_imread(file_path):
    cv_img = cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), -1)
    return cv_img
# opencv写入中文路径名会乱码
def cv_imwrite(img, file_path):
    cv2.imencode('.jpg', img)[1].tofile(file_path)
# 把原图裁剪为多个小图(900*383)
def crop_little_pic(pic_path):
    img = cv_imread(pic_path)
    (sh, sw) = img.shape[:2]
    # 将图片的宽设置为900，高则按比例缩放
    res = cv2.resize(img, (900, round(sh * (900 / sw))))
    # 获取缩放后的高和宽，判断图片可裁剪的张数
    (ch, cw) = res.shape[:2]
    crop_pic_count = int(ch / 383)
    # 计算Y轴偏移
    start_y = int(ch % 383 / 2)
    # 根据图片的张数来决定怎么裁剪
    for i in range(0, crop_pic_count):
        crop_img = res[383 * i + start_y: 383 * (i + 1) + start_y, 0:900]
        cv_imwrite(crop_img, os.path.join(pic_crop_dir, str(int(round(time.time() * 1000))) + '.jpg'))
# 绘制文字
def draw_text(pic_path, date):
    img = cv_imread(pic_path)
    # 将图片从OpenCv格式转为PIL格式
    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 加载字体（字体文件名，字体大小）
    title_font = ImageFont.truetype('apple-simple.ttf', 52)
    date_font = ImageFont.truetype('apple-simple.ttf', 44)
    # 绘制文字的位置
    title_pos = (236, 110)
    date_pos = (316, 192)
    # 绘制内容
    title_content = u"『抠腚早报速读』"
    date_content = u"第%s期" % date[2:]
    # 绘制
    draw = ImageDraw.Draw(img_pil)
    draw.text(title_pos, title_content, font=title_font, fill=(255, 255, 255))
    draw.text(date_pos, date_content, font=date_font, fill=(255, 255, 255))
    img_open_cv = cv2.cvtColor(np.asarray(img_pil), cv2.COLOR_RGB2BGR)
    cv_imwrite(img_open_cv, os.path.join(pic_font_dir, date[2:] + '.jpg'))
# 遍历获得某类文件路径列表
def fetch_file_path(path, file_type):
    file_list = []
    f = os.listdir(path)
    for i in f:
        if i.endswith(file_type):
            file_list.append(os.path.join(path, i))
    return file_list
# 构造生成日期列表
def init_date_list(begin_date, count):
    d_list = []
    begin_date = datetime.strptime(begin_date, "%Y%m%d")
    end_date = datetime.strptime((datetime.now() + timedelta(days=count)).strftime("%Y%m%d"), "%Y%m%d")
    while begin_date <= end_date:
        date_str = begin_date.strftime("%Y%m%d")
        d_list.append(date_str)
        begin_date += timedelta(days=1)
    return d_list
if __name__ == '__main__':
    is_dir_existed(pic_source_dir)
    while True:
        choice = input(
            "%s\n请输入你想进行的操作\n1.进行图片裁剪\n2.图片加字\n3.清空裁剪文件夹\n4.清空加字文件夹\n5.退出程序\n%s\n" % ('=' * 32, '=' * 32))
        if choice == '1':
            is_dir_existed(pic_crop_dir)
            pic_path_list = fetch_file_path(pic_source_dir, ".jpg")
            if len(pic_path_list) == 0:
                print("原图文件夹中无图片，请先添加图片！")
            else:
                print("开始批量裁剪...")
                begin = datetime.now()
                for pic in pic_path_list:
                    crop_little_pic(pic)
                end = datetime.now()
                print("批量裁剪完毕，生成图片：%d张，耗时：%s秒" % (len(fetch_file_path(pic_crop_dir, ".jpg")), (end - begin).seconds))
        elif choice == '2':
            is_dir_existed(pic_font_dir)
            crop_path_list = fetch_file_path(pic_crop_dir, ".jpg")
            date_list = init_date_list(start_date, len(crop_path_list) + 1)
            if len(crop_path_list) == 0:
                print("裁剪文件夹中无图片，请先生成裁剪图片！")
            else:
                print("开始批量加字...")
                begin = datetime.now()
                print(len(crop_path_list), len(date_list))
                for i in range(len(crop_path_list)):
                    draw_text(crop_path_list[i], date_list[i])
                end = datetime.now()
                print("批量加字完毕，处理图片：%d张，耗时：%s秒" % (len(fetch_file_path(pic_font_dir, ".jpg")), (end - begin).seconds))
        elif choice == '3':
            if is_dir_existed(pic_crop_dir, False):
                shutil.rmtree(pic_crop_dir)
                print("文件夹删除成功！")
            else:
                print("文件夹不存在，删除失败~")
        elif choice == '4':
            if is_dir_existed(pic_font_dir, False):
                shutil.rmtree(pic_font_dir)
                print("文件夹删除成功！")
            else:
                print("文件夹不存在，删除失败~")
        elif choice == '5':
            exit("退出程序~")
        else:
            print("错误序号，请确认后重新输入!!!")

# 截取视频的第一帧
def fetch_video_first_frame(path_list):
    for mp4 in path_list:
        cap = cv2.VideoCapture(mp4)
        if cap.isOpened():
            ret, im = cap.read()
            cv2.imencode('.jpg', im)[1].tofile(
                os.path.join(pic_source_output_dir, mp4.split("\\")[-1]).replace("mp4", "jpg"))
        cap.release()