package main

import (
	"flag"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	_ "image/png"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/nfnt/resize"
)

/*
1. 读取验证码
2. 将验证码图片二值化
3. 将二值化后的图片进行近似等分切割(可能有更好的方法,对于精度不高要求不高的话等分即可)
4. 将切割后的图片进行识别(调用train()函数,手动标注验证码,生成指纹数组,也可以写入到一个文件中,然后读取)
5. 将识别后的结果进行拼接
*/
var (
	model               = flag.String("m", "train", "运行模式【train/predict】")
	imageFile           = flag.String("img", "./demo.jpg", "要识别的图像文件")
	threshold           = flag.Uint64("ct", 0xcccc, "颜色过滤阈值")
	clearColorThreshold uint32
)

func init() {
	flag.Parse()
	clearColorThreshold = uint32(*threshold)
	if _, err := os.Stat("./view"); err != nil {
		os.Mkdir("./view", 0777)
	}
}

// 判断r四周是否有任意两个连通白块，如果没有，则认为是噪点
func imgConnect(ru, rd, rl, rr uint32) bool {
	if (ru < clearColorThreshold &&
		(rl < clearColorThreshold || rr < clearColorThreshold)) ||
		(rd < clearColorThreshold && (rl < clearColorThreshold || rr < clearColorThreshold)) ||
		(rl < clearColorThreshold && (ru < clearColorThreshold || rd < clearColorThreshold)) ||
		(rr < clearColorThreshold && (ru < clearColorThreshold || rd < clearColorThreshold)) {
		return true
	}
	return false
}

// 二值化
func imgBarbarization(img image.Image) image.Image {
	binImg := image.NewGray16(img.Bounds())
	draw.Draw(binImg, binImg.Bounds(), img, img.Bounds().Min, draw.Over)
	rect := binImg.Bounds()
	// 遍历点像素点
	for x := 0; x < rect.Dx(); x++ {
		for y := 0; y < rect.Dy(); y++ {
			// 获取颜色
			r, _, _, _ := binImg.At(x, y).RGBA()
			ru, _, _, _ := img.At(x, y+1).RGBA()
			rd, _, _, _ := img.At(x, y-1).RGBA()
			rl, _, _, _ := img.At(x-1, y).RGBA()
			rr, _, _, _ := img.At(x+1, y).RGBA()
			if r < clearColorThreshold && imgConnect(ru, rd, rl, rr) {
				binImg.Set(x, y, color.White)
			} else {
				binImg.Set(x, y, color.Black)
			}
		}
	}
	return binImg
}

// 删除图片无用外部轮廓
func imgCutSide(img image.Image) image.Image {
	rect := img.Bounds()
	// 开始x坐标
	leftStartX := 0
	for x := 0; x < rect.Dx(); x++ {
		lxflag := false
		for y := 0; y < rect.Dy(); y++ {
			// 获取颜色
			r, _, _, _ := img.At(x, y).RGBA()
			if r == 0xFFFF {
				lxflag = true
				break
			}
		}
		if lxflag {
			leftStartX = x
			break
		}
	}
	// 开始y坐标
	leftStartY := 0
	for y := 0; y < rect.Dy(); y++ {
		lyflag := false
		for x := 0; x < rect.Dx(); x++ {
			// 获取颜色
			r, _, _, _ := img.At(x, y).RGBA()
			if r == 0xFFFF {
				lyflag = true
				break
			}
		}
		if lyflag {
			leftStartY = y
			break
		}
	}
	// 结束x坐标
	rightEndX := 0
	for x := rect.Dx(); x > 0; x-- {
		rxflag := false
		for y := rect.Dy(); y > 0; y-- {
			// 获取颜色
			r, _, _, _ := img.At(x, y).RGBA()
			if r == 0xFFFF {
				rxflag = true
				break
			}
		}
		if rxflag {
			rightEndX = x
			break
		}
	}
	// 结束y坐标
	rightEndY := 0
	for y := rect.Dy(); y > 0; y-- {
		ryflag := false
		for x := rect.Dy(); x > 0; x-- {
			// 获取颜色
			r, _, _, _ := img.At(x, y).RGBA()
			if r == 0xFFFF {
				ryflag = true
				break
			}
		}
		if ryflag {
			rightEndY = y
			break
		}
	}
	// 创建新的图片
	rectangle := image.Rectangle{
		Min: image.Point{X: 0, Y: 0},
		Max: image.Point{X: rightEndX - leftStartX, Y: rightEndY - leftStartY},
	}
	newSideImg := image.NewGray16(rectangle)
	draw.Draw(newSideImg, newSideImg.Bounds(), img, image.Point{X: leftStartX, Y: leftStartY}, draw.Over)
	return newSideImg
}

func cutImage(img image.Image) []image.Image {
	rect := img.Bounds()
	imgs := make([]image.Image, 0)
	// 记录当前x的位置
	nowCutStartX := 0
	for x := 0; x <= rect.Dx(); x++ {
		every := rect.Dx() / 4
		if x%every == 0 && x != 0 {

			// 创建新的图片
			rectangle := image.Rectangle{
				Min: image.Point{X: 0, Y: 0},
				Max: image.Point{X: x + 3 - nowCutStartX, Y: rect.Dy()},
			}
			newImg := image.NewGray16(rectangle)
			draw.Draw(newImg, newImg.Bounds(), img, image.Point{X: nowCutStartX, Y: 0}, draw.Over)
			newResizeImg := imgCutSide(newImg)
			newResizeImg = resize.Resize(16, 16, newResizeImg, resize.Lanczos3)
			newResizeImg = imgCutSide(newImg)
			newResizeImg = resize.Resize(16, 16, newResizeImg, resize.Lanczos3)
			imgs = append(imgs, newResizeImg)
			nowCutStartX = x + 3
		}
	}

	return imgs
}

// 指纹验证
func imgSign(imgBinary string) string {
	if imgBinary == "" {
		return imgBinary
	}
	imgBinarySign := strings.Split(imgBinary, "")
	// 相似度
	maxSimilarity := 0
	num := ""
	for _, sign := range numSign {
		tmpSimilarity := 0
		signArr := strings.Split(sign[1], "")
		for k, s := range imgBinarySign {
			if s == signArr[k] {
				tmpSimilarity++
			}
		}
		if maxSimilarity < tmpSimilarity {
			maxSimilarity = tmpSimilarity
			num = sign[0]
		}
	}
	return num
}

// 将切割后的字符图片转换为二进制, 与指纹库进行比对
func recognition(imgs []image.Image) (string, []string) {
	signNum := make([]string, 0)
	for _, img := range imgs {
		tmpSignNum := make([]string, 0)
		rect := img.Bounds()
		for x := 0; x < rect.Dx(); x++ {
			for y := 0; y < rect.Dy(); y++ {
				// 获取颜色
				r, _, _, _ := img.At(x, y).RGBA()
				if r > 0x7788 {
					tmpSignNum = append(tmpSignNum, "1")
				} else {
					tmpSignNum = append(tmpSignNum, "0")
				}
			}
		}
		signNum = append(signNum, strings.Join(tmpSignNum, ""))
	}
	res := make([]string, 0)
	for _, v := range signNum {
		res = append(res, imgSign(v))
	}
	return strings.Join(res, ""), signNum
}

func train() {
	stringgen := ""
	f, err := os.Create("train.txt")
	if err != nil {
		fmt.Println(err)
		return
	}
	client := &http.Client{
		Timeout: 2 * time.Second,
	}
	for {
		resp, err := client.Get("https://beian.china-eia.com/servlet/validateCodeServlet?" + strconv.FormatInt(time.Now().UnixNano()/1e6, 10))
		if err != nil {
			fmt.Sprintf("打开图片失败:%+v", err.Error())
			continue
		}
		fmt.Println(time.Now().UnixNano() / 1e6)

		img, err := jpeg.Decode(resp.Body)
		if err != nil {
			fmt.Sprintf("解析图片失败:%+v", err.Error())
			continue
		}
		resp.Body.Close()
		dts, err := os.Create("./view/01.原始图片.jpeg")
		jpeg.Encode(dts, img, nil)
		// 图片二值化
		binImg := imgBarbarization(img)
		dts1, err := os.Create("./view/02.二值化后的图片.jpeg")
		jpeg.Encode(dts1, binImg, nil)
		// 图片切割
		imgCutSilce := cutImage(binImg)
		// 输出切割后的四张图片
		dts3, err := os.Create("./view/03.第一个字符.jpeg")
		jpeg.Encode(dts3, imgCutSilce[0], nil)
		dts4, err := os.Create("./view/04.第二个字符.jpeg")
		jpeg.Encode(dts4, imgCutSilce[1], nil)
		dts5, err := os.Create("./view/05.第三个字符.jpeg")
		jpeg.Encode(dts5, imgCutSilce[2], nil)
		dts6, err := os.Create("./view/06.第四个字符.jpeg")
		jpeg.Encode(dts6, imgCutSilce[3], nil)
		// 验证码识别
		result, resu := recognition(imgCutSilce)
		fmt.Println("程序识别结果: ", result)
		x := ""
		fmt.Print("请输入./view/01.原始图片正确结果:  ")
		fmt.Scan(&x)
		if x == "q" {
			return
		} else if x == "k" {
			continue
		} else if len(x) != 4 {
			fmt.Println("输入的格式不对")
			continue
		}

		inputs := strings.Split(x, "")
		stringgen = fmt.Sprintf("{\"%s\", \"%s\"},\n{\"%s\", \"%s\"},\n{\"%s\", \"%s\"},\n{\"%s\",  \"%s\"},\n", inputs[0], resu[0], inputs[1], resu[1], inputs[2], resu[2], inputs[3], resu[3])

		_, err = f.Write([]byte(stringgen))
		if err != nil {
			fmt.Println(err)
		}
	}
}

// 识别预测
func predict(imageFile string) (string, error) {
	fi, err := os.Open(imageFile)
	if err != nil {
		return "", err
	}
	defer fi.Close()
	img, err := jpeg.Decode(fi)
	if err != nil {
		return "", err
	}
	// 图片二值化
	binImg := imgBarbarization(img)
	// 图片切割
	imgCutSilce := cutImage(binImg)
	// 验证码识别
	result, resu := recognition(imgCutSilce)
	fmt.Println("程序识别结果: ", result, resu)
	return result, nil
}

func main() {
	switch *model {
	case "train":
		train()
	default:
		predict(*imageFile)
	}
}