import cv2 as cv
import matplotlib.pyplot as plt
import numpy as np
import copy
import os
import math
import pandas as pd
from sklearn.cluster import  KMeans

def rgb2gray(rgb):
    fig = np.dot(rgb[...,:3],[0.229,0.587,0.114])
    plt.imshow(fig, cmap='gray')
    plt.axis('off')
    # plt.title("gray")
    plt.show()
    return fig

def valid_cov(res):
    new_fig = copy.deepcopy(res)

    width = res.shape[0]
    height = res.shape[1]
    block_size = 50
    w_block = math.ceil(width/block_size)
    h_block = math.ceil(height/block_size)
    for x in range(w_block):
        for y in range(h_block):
            x_ind = x*block_size
            y_ind = y*block_size
            cnt = 0
            val = 0
            for xi in range(x_ind, x_ind+block_size):
                if xi >= width:
                    break
                for yi in range(y_ind, y_ind+block_size):
                    if yi >= height:
                        break
                    cnt += 1
                    val += new_fig[xi, yi]
            avg = val / cnt
            #print('abg=', avg)
            if avg < 85:
                for xi in range(x_ind, x_ind + block_size):
                    if xi >= width:
                        break
                    for yi in range(y_ind, y_ind + block_size):
                        if yi >= height:
                            break
                        new_fig[xi, yi] = 0

            else:
                for xi in range(x_ind, x_ind + block_size):
                    if xi >= width:
                        break
                    for yi in range(y_ind, y_ind + block_size):
                        if yi >= height:
                            break
                        new_fig[xi, yi] = 255

    for x in range(w_block):
        if x == 0 or (x+1)*block_size >= width:
            break
        for y in range(h_block):
            if y == 0 or (y+1)*block_size >= height:
                break
            cnt = 0
            if new_fig[(x-1)*block_size, y*block_size] == 255:
                cnt += 1
            if new_fig[(x+1)*block_size, y*block_size] == 255:
                cnt += 1
            if new_fig[x*block_size, (y-1)*block_size] == 255:
                cnt += 1
            if new_fig[x*block_size, (y+1)*block_size] == 255:
                cnt += 1
            if cnt >=3:
                print('start remove x=', x, ' y=', y)
                for xi in range(x*block_size, (x+1)*block_size):
                    if xi >= width:
                        break
                    for yi in range(y*block_size, (y+1)*block_size):
                        if yi >= height:
                            break
                        new_fig[xi, yi] = 255

    plt.imshow(new_fig, cmap='gray')
    plt.axis('off')
    plt.show()
    # plt.title("valid cov")
    return new_fig

def one_setp_sub(res):
    new_fig = copy.deepcopy(res)
    width = res.shape[0]
    height = res.shape[1]
    for x in range(width):
        for y in range(height):
            if y+1 < height:
                t = res[x, y] - res[x, y + 1]
                if t > 1:
                    new_fig[x, y] = 255
                else:
                    new_fig[x, y] = 0
    plt.imshow(new_fig, cmap='gray')
    plt.axis('off')
    plt.show()
    # plt.title("one step sub")
    return new_fig

def dim3to2(rgb):
    const_val = math.sqrt(2)/2
    return np.dot(rgb[...,:3],np.array([[0,const_val,-const_val], [const_val,-const_val,0]]).T)

def count(dim2res, img):
    width = dim2res.shape[0]
    height = dim2res.shape[1]
    green_yellow = copy.deepcopy(img)

    cov_cnt = 0#统计岩石面积
    cac_pct = 0#统计绿色和黄色
    for x in range(width):
        for y in range(height):
            if not rock[x, y] == 255:
                green_yellow[x, y, :] = [0, 0, 0]
                continue
            cov_cnt += 1
            if dim2res[x, y, 0]**2 + dim2res[x, y, 1]**2 > 30**2 and dim2res[x, y, 0] > dim2res[x, y, 1]:
                green_yellow[x, y, :] = [0, 255, 0]
                cac_pct += 1

            else:
                green_yellow[x, y, :] = [255, 255, 255]

    plt.imshow(green_yellow)
    plt.axis('off')
    # plt.title("light cov")
    plt.show()
    return green_yellow, cov_cnt, cac_pct

#------------------main---------------------------------------------------------
path = 'D:\\temp\\B题全部数据\\Rock'
filenames = os.listdir(path)
pct_list = pd.DataFrame(columns=['file name', 'oil percentage'])
for filename in filenames:
    sp_list = filename.split('.')
    if not sp_list[-1] == 'jpg':
        continue

    first = sp_list[0]#1-2
    type = first.split('-')[1]#2
    if type != '2':#只有2是荧光图
        continue

    fig = plt.imread(path + '\\' + filename)
    b, g, r = cv.split(fig)
    img = cv.merge([r, g, b])  # 为了适应fc中元组顺序为RGB，故调整
    res = rgb2gray(img)

    #一阶差分
    new_fig = one_setp_sub(res)

    #计算有效面积
    rock = valid_cov(new_fig)

    #计算荧光
    b, g, r = cv.split(fig)
    img = cv.merge([r, g, b])  # 为了适应fc中元组顺序为RGB，故调整
    dim2res = dim3to2(img)


    green_yellow, cov_cnt, cac_pct = count(dim2res, img)


    percent = cac_pct / cov_cnt
    pct_list.append({'file name': filename, 'oil percentage': percent}, ignore_index=True)
    print('filename=', filename, ' percentage=', percent)


pct_list.to_csv('D:\\其他\\data\\rock_oil_jpg.csv')