# coding=utf-8
__author__ = 'vivienfanghua'

import numpy as np
import math
import random
# from PIL import Image
import matplotlib.pyplot as plt


class Model:
    def __init__(self, size):
        # 常量
        self.size = int(size)
        self.vCH4 = 1  # 初始浓度
        self.vSO4 = 10  # 初始浓度
        self.vH2 = 1e-7  # 初始浓度
        self.vH2S = 1e-10  # 初始浓度
        self.m = 0  # 维持能

        # 变量
        self.CH4 = np.ones((self.size, self.size)).dot(self.vCH4)
        self.H2 = np.ones((self.size, self.size)).dot(self.vH2)
        self.SO4 = np.ones((self.size, self.size)).dot(self.vSO4)
        self.ANME = np.zeros((self.size, self.size))
        self.SRB = np.zeros((self.size, self.size))
        self.GFE1 = np.zeros((self.size, self.size))
        self.GFE2 = np.zeros((self.size, self.size))
        self.DB = np.zeros((self.size, self.size))  # 死亡生物量
        self.PG = np.zeros((self.size, self.size))  # 公共物质
        self.vANME = np.zeros((self.size, self.size))  # ANME生长速率
        self.vSRB = np.zeros((self.size, self.size))  # SRB生长速率
        self.oANME = 150  # ANME初始生物量
        self.oSRB = 150  # SRB初始生物量
        self.sANME = 300  # ANME生物量分裂阈值
        self.sSRB = 300  # SRB生物量分裂阈值
        self.dANME = 15  # ANME生物量生存阈值
        self.dSRB = 15  # SRB生物量生存阈值

        self.af = 0.1

    # max_iterations:迭代次数
    # interval:间隔多少次输出图像
    def iteration(self, max_iterations, interval):
        for i in range(max_iterations):
            print("Iteration " + str(i))
            if i % interval == 0:
                rgbarray = np.zeros((self.size, self.size, 3), 'uint8')
                # r_channel = np.ones(self.ANME.shape) * 200 * (self.ANME>0)
                # b_channel = np.ones(self.SRB.shape) * 200 * (self.SRB > 0)
                np.savetxt("output/srb_" + str(i), self.SRB, fmt="%.8e")
                np.savetxt("output/anme_" + str(i), self.ANME, fmt="%.8e")
                # np.savetxt("output/ch4_"+str(i),self.CH4,fmt="%.8e")
                np.savetxt("output/so4_" + str(i), self.SO4, fmt="%.8e")
                np.savetxt("output/h2_" + str(i), self.H2, fmt="%.8e")
                np.savetxt("output/gfe1_" + str(i), self.GFE1, fmt="%.8e")
                np.savetxt("output/gfe2_" + str(i), self.GFE2, fmt="%.8e")
                np.savetxt("output/vsrb_" + str(i), self.vSRB, fmt="%.8e")
                np.savetxt("output/vanme_" + str(i), self.vANME, fmt="%.8e")

                rgbarray[..., 0] = self.ANME / self.sANME * 255  # R红色通道 TODO>300
                rgbarray[..., 1] = self.SRB / self.sSRB * 255  # B蓝色通道 TODO>300
                rgbarray[self.ANME > 300, 0] = 255
                rgbarray[self.SRB > 300, 1] = 255
                # rgbarray[..., 0] = 255
                # rgbarray[..., 0] = r_channel  # R红色通道
                # rgbarray[..., 2] = b_channel  # B蓝色通道
                # img = Image.fromarray(rgbarray)
                # img.save('merge_' + str(i) + '.jpg')
                plt.imshow(rgbarray)
                plt.show()
                # anme_img = Image.fromarray(self.ANME)
                # anme_img = anme_img.convert('L')
                # anme_img.save('ANME_' + str(i) + '.jpg')
                # srb_img = Image.fromarray(self.SRB)
                # srb_img = srb_img.convert('L')
                # srb_img.save('SRB_' + str(i) + '.jpg')

            # 根据物质和细胞设置不同间隔
            self.update_CH4()
            self.update_GFE1()
            self.update_vANME()
            self.update_ANME()
            self.update_H21()
            count = 1
            while (count < 3):
                self.update_H2()
                count += 1
            self.update_GFE2()
            self.update_vSRB()
            self.update_SRB()
            self.update_H22()
            self.update_SO4()
            count = 1
            while (count < 100):
                self.update_H2()
                count += 1
            count = 1
            while (count < 100):
                self.update_SO4k()
                count += 1

    # 初始化生物状态
    def init_para(self):

        for (x, y), value in np.ndenumerate(self.ANME):
            if (x + y) % 2:
                self.ANME[x, y] = 0
            else:
                self.ANME[x, y] = self.oANME
        for (x, y), value in np.ndenumerate(self.SRB):
            if (x + y) % 2:
                self.SRB[x, y] = self.oSRB
            else:
                self.SRB[x, y] = 0

    # 初始化方法1
    def init_para1(self, width):
        start_x = random.randint(0, self.size - width)
        start_y = random.randint(0, self.size - width)
        for i in range(start_x, start_x + width):
            for j in range(start_y, start_y + width):
                if (i + j) % 2 == 0:
                    self.ANME[i, j] = self.oANME
                else:
                    self.SRB[i, j] = self.oSRB

    # 随机产生num个点
    def init_para2(self, num):
        idx = random.sample(range(self.size * self.size), num)
        for i in range(len(idx)):
            if i % 2:
                self.ANME[divmod(idx[i], self.size)] = self.oANME
            else:
                self.SRB[divmod(idx[i], self.size)] = self.oSRB

    # 更新CH4
    def update_CH4(self):
        # 边界补齐
        self.CH4[0:1] = self.vCH4
        self.CH4[self.size - 1:self.size] = self.vCH4
        self.CH4[0:self.size, 0:1] = self.vCH4
        self.CH4[0:self.size, self.size - 1:self.size] = self.vCH4
        # 物质消耗
        self.CH4 += self.vANME * self.ANME * -10.12 - 0.23 * self.af * 0.11 / 365 * self.ANME  # 公式参数理解
        # 物质扩散
        self.CH4 = self.diffuse(self.CH4)
        # 物质补齐
        for (x, y), value in np.ndenumerate(self.CH4):
            if self.CH4[x, y] < 0:
                self.CH4[x, y] = 5e-10

    # 扩散CH4
    def update_CH4k(self):
        # 边界补齐
        self.CH4[0:1] = self.vCH4
        self.CH4[self.size - 1:self.size] = self.vCH4
        self.CH4[0:self.size, 0:1] = self.vCH4
        self.CH4[0:self.size, self.size - 1:self.size] = self.vCH4
        # 物质扩散
        self.CH4 = self.diffuse(self.CH4)
        # 物质补齐
        for (x, y), value in np.ndenumerate(self.CH4):
            if self.CH4[x, y] < 0:
                self.CH4[x, y] = 5e-10

    # 消耗SO4
    def update_SO4(self):
        # 边界补齐
        self.SO4[0:1] = self.vSO4
        self.SO4[self.size - 1:self.size] = self.vSO4
        self.SO4[0:self.size, 0:1] = self.vSO4
        self.SO4[0:self.size, self.size - 1:self.size] = self.vSO4
        # 物质消耗
        self.SO4 += self.vSRB * self.SRB * -41.2  # 公式理解
        # self.SO4 = self.SO4 - self.SRB.dot(self.cSO4)
        # 物质扩散
        self.SO4 = self.diffuse(self.SO4)
        # 物质补齐
        for (x, y), value in np.ndenumerate(self.SO4):
            if self.SO4[x, y] < 0:
                self.SO4[x, y] = 5e-10

    # 扩散SO4
    def update_SO4k(self):
        # 边界补齐
        self.SO4[0:1] = self.vSO4
        self.SO4[self.size - 1:self.size] = self.vSO4
        self.SO4[0:self.size, 0:1] = self.vSO4
        self.SO4[0:self.size, self.size - 1:self.size] = self.vSO4
        # 物质扩散
        self.SO4 = self.diffuse(self.SO4)
        # 物质补齐
        for (x, y), value in np.ndenumerate(self.SO4):
            if self.SO4[x, y] < 0:
                self.SO4[x, y] = 5e-10

    # 更新产生H2
    def update_H21(self):
        for (x, y), value in np.ndenumerate(self.H2):
            if (self.ANME[x, y] > 0 and self.vANME[x, y] and self.GFE1[x, y] < 0):
                # 产生H2
                self.H2[x, y] += self.vANME[x, y] * self.ANME[x, y] * 4.88  # 公式理解
            if self.H2[x, y] < 0:
                self.H2[x, y] = 0
                # 物质扩散
                # self.H2 = self.diffuse(self.H2)
                # 边界补齐
                # self.H2[0:1] = self.vH2
                # self.H2[self.size - 1:self.size] = self.vH2
                # self.H2[0:self.size, 0:1] = self.vH2
                # self.H2[0:self.size, self.size - 1:self.size] = self.vH2


                # 更新消耗H2

    def update_H22(self):
        for (x, y), value in np.ndenumerate(self.H2):
            if (self.SRB[x, y] > 0 and self.vSRB[x, y] > 0 and self.GFE2[x, y] < 0):
                # 消耗H2
                self.H2[x, y] += self.vSRB[x, y] * self.SRB[x, y] * (-3.6)  # 公式理解
            if self.H2[x, y] < 0:
                self.H2[x, y] = self.vH2
                # 物质扩散
                # self.H2 = self.diffuse(self.H2)
                # 边界补齐
                # self.H2[0:1] = self.vH2
                # self.H2[self.size - 1:self.size] = self.vH2
                # self.H2[0:self.size, 0:1] = self.vH2
                # self.H2[0:self.size, self.size - 1:self.size] = self.vH2

    # 更新扩散H2
    def update_H2(self):
        # 物质扩散
        self.H2 = self.diffuse(self.H2)
        # 边界补齐
        self.H2[0:1] = self.vH2
        self.H2[self.size - 1:self.size] = self.vH2
        self.H2[0:self.size, 0:1] = self.vH2
        self.H2[0:self.size, self.size - 1:self.size] = self.vH2

    # 物质扩散通用方法
    def diffuse(self, matrix):
        change = np.zeros(matrix.shape)
        for (x, y), value in np.ndenumerate(matrix):
            v = matrix[x, y]
            if y < self.size - 1: change[x, y] += 0.2 * (matrix[x, y + 1] - v)
            if x < self.size - 1: change[x, y] += 0.2 * (matrix[x + 1, y] - v)
            if x > 0: change[x, y] += 0.2 * (matrix[x - 1, y] - v)
            if y < self.size - 1: change[x, y] += 0.2 * (matrix[x, y + 1] - v)
            if x < self.size - 1 and y < self.size - 1: change[x, y] += 0.05 * (matrix[x + 1, y + 1] - v)
            if x < self.size - 1 and y > 0: change[x, y] += 0.05 * (matrix[x + 1, y - 1] - v)
            if x > 0 and y > 0: change[x, y] += 0.05 * (matrix[x - 1, y - 1] - v)
            if x > 0 and y < self.size - 1: change[x, y] += 0.05 * (matrix[x - 1, y + 1] - v)
        return matrix + change * 0.3  # 扩散系数0.1

    # 吉布斯自由能1
    def update_GFE1(self):
        for (x, y), value in np.ndenumerate(self.GFE1):
            if self.ANME[x, y]:
                self.GFE1[x, y] = 16.92 + 8.314e-3 * 298 * math.log(self.H2[x, y] ** 0.5 / self.CH4[x, y] ** 0.125)
            else:
                self.GFE1[x, y] = 0

    # 吉布斯自由能2
    def update_GFE2(self):
        for (x, y), value in np.ndenumerate(self.GFE2):
            if self.SRB[x, y]:
                self.GFE2[x, y] = -19.04 + 8.314e-3 * 298 * math.log(
                    self.vH2S ** 0.125 / ((self.H2[x, y] ** 0.5) * (self.SO4[x, y] ** 0.125)))
            else:
                self.GFE2[x, y] = 0

    # ANME生长速率
    def update_vANME(self):
        v0 = 0.1 / 24 / 60
        for (x, y), value in np.ndenumerate(self.vANME):
            if self.ANME[x, y]:
                fk = self.CH4[x, y] / (1.5e-3 + self.CH4[x, y])
                ft = 1 - math.exp((self.GFE1[x, y] + 0) / 8.314e-3 / 298)
                self.vANME[x, y] = v0 * fk * ft
            else:
                self.vANME[x, y] = 0

    # SRB生长速率
    def update_vSRB(self):
        v0 = 0.29 / 24 / 60
        for (x, y), value in np.ndenumerate(self.vSRB):
            if self.SRB[x, y]:
                # fk = self.H2[x, y] * self.SO4[x, y] / (1.5e-3 + self.H2[x, y]) / (1.5e-3 + self.SO4[x, y])
                ft = 1 - math.exp((self.GFE2[x, y] + 0) / 8.314e-3 / 298)
                fk = 0.65
                self.vSRB[x, y] = v0 * ft * fk
            else:
                self.vSRB[x, y] = 0

    # 更新ANME状态
    def update_ANME(self):
        for (x, y), value in np.ndenumerate(self.ANME):
            # 生物量累积
            if self.GFE1[x, y] < 0:
                self.ANME += 100 * (self.vANME - 0.11 / 365 / 24 / 60) * self.ANME  # 负增长是否对应负速率？吉布斯自由能是否还需要判断？
                # 分裂
                if self.ANME[x, y] > self.sANME:
                    neighbor = [(x - 1, y - 1), (x - 1, y), (x - 1, y + 1), (x, y - 1), (x, y + 1), (x + 1, y - 1),
                                (x + 1, y), (x + 1, y + 1)]
                    # if x == 0 or x == self.size - 1 or y == 0 or y == self.size - 1:
                    # neighbor = [i for i in neighbor if self.size > i[0] >= 0 and self.size > i[1] >= 0]  # 边界判断
                    # neighbor = [i for i in neighbor if
                    #             self.size > i[0] >= 0 and self.size > i[1] >= 0 and self.ANME[i[0], i[1]] == 0 and self.SRB[
                    #                 i[0], i[1]] == 0]
                    self.ANME[x, y] *= 0.5  # TODO 分裂后生物量如何变？
                    neighbor_t = [i for i in neighbor if
                                  self.size > i[0] >= 0 and self.size > i[1] >= 0]  # 有效的八邻域
                    neighbor_a = [i for i in neighbor_t if
                                  self.ANME[i[0], i[1]] == 0 and self.SRB[i[0], i[1]] == 0]  # 带空位的八邻域
                    if neighbor_a:
                        rand_pos = neighbor_a[random.randint(0, len(neighbor_a) - 1)]  # 随机选择一个位置分裂
                        # self.ANME[x, y] *= 0.5  # TODO 分裂后生物量如何变？
                        self.ANME[rand_pos] = self.ANME[x, y]
                    else:
                        rand_pos = neighbor_t[random.randint(0, len(neighbor_t) - 1)]  # 随机选择一个位置分裂
                        if self.ANME[rand_pos]:
                            self.ANME[rand_pos] = [self.ANME[x, y], self.ANME[rand_pos]][
                                random.randint(0, 1)]  # 占位:二分之一概率选择
                        elif self.SRB[rand_pos]:
                            if random.randint(0, 1):
                                self.ANME[rand_pos] = self.ANME[x, y]
                                self.SRB[rand_pos] = 0
                            else:
                                self.ANME[rand_pos] = 0
                if self.ANME[x, y] < self.dANME:
                    self.ANME[x, y] = 0  # ANME死亡 TODO 能量和物质如何释放？

                    # 更新ANME状态

    def update_SRB(self):
        for (x, y), value in np.ndenumerate(self.SRB):
            # 生物量累积
            if self.GFE2[x, y] < 0:
                self.SRB += 100 * (self.vSRB - 0.11 / 365 / 24 / 60) * self.SRB  # 负增长是否对应负速率？吉布斯自由能是否还需要判断？
                # 分裂
                if self.SRB[x, y] > self.sSRB:
                    self.SRB[x, y] *= 0.5  # TODO 分裂后生物量如何变？
                    neighbor = [(x - 1, y - 1), (x - 1, y), (x - 1, y + 1), (x, y - 1), (x, y + 1), (x + 1, y - 1),
                                (x + 1, y), (x + 1, y + 1)]
                    # if x == 0 or x == self.size - 1 or y == 0 or y == self.size - 1:
                    # neighbor = [i for i in neighbor if self.size > i[0] >= 0 and self.size > i[1] >= 0]  # 边界判断
                    # rand_pos = neighbor[random.randint(0, len(neighbor) - 1)]  # 随机选择一个位置分裂
                    neighbor_t = [i for i in neighbor if
                                  self.size > i[0] >= 0 and self.size > i[1] >= 0]  # 有效的八邻域
                    neighbor_a = [i for i in neighbor_t if
                                  self.ANME[i[0], i[1]] == 0 and self.SRB[i[0], i[1]] == 0]  # 带空位的八邻域
                    if neighbor_a:  # 有空位情况
                        rand_pos = neighbor_a[random.randint(0, len(neighbor_a) - 1)]  # 随机选择一个位置分裂

                        self.SRB[rand_pos] = self.SRB[x, y]
                    else:  # 没有空位情况
                        rand_pos = neighbor_t[random.randint(0, len(neighbor_t) - 1)]
                        if self.SRB[rand_pos]:
                            self.SRB[rand_pos] = [self.SRB[x, y], self.SRB[rand_pos]][
                                random.randint(0, 1)]  # 占位:二分之一概率选择
                        elif self.ANME[rand_pos]:
                            if random.randint(0, 1):
                                self.SRB[rand_pos] = self.SRB[x, y]
                                self.ANME[rand_pos] = 0  # ANME 竞争失败
                            else:
                                self.SRB[rand_pos] = 0  # ANME竞争成功
                if self.SRB[x, y] < self.dSRB:
                    self.SRB[x, y] = 0  # SRB死亡 TODO 能量和物质如何释放？


if __name__ == "__main__":
    m = Model(50)
    # m.init_para2(8)

    m.init_para1(4)
    m.iteration(100, 10)