# 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 = 1  # 初始浓度
        self.vH2 = 5e-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.0  # ANME初始生物量
        self.oSRB = 150.0  # SRB初始生物量
        self.sANME = 300.0  # ANME生物量分裂阈值
        self.sSRB = 300.0  # SRB生物量分裂阈值
        self.dANME = 30.0  # ANME生物量生存阈值
        self.dSRB = 30.0  # 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')
                rgbarray[..., 0] = self.ANME  # R红色通道
                rgbarray[..., 2] = self.SRB  # 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_SO4()
            self.update_H2()
            self.update_GFE2()
            self.update_vSRB()
            self.update_SRB()
            self.update_H2()
            self.update_GFE1()
            self.update_vANME()
            self.update_ANME()
            self.update_H2()

    # 初始化生物状态
    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)

    # 更新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)

    # 更新H2
    def update_H2(self):
        # 产生H2
        self.H2 += self.vANME * self.ANME * 4.88  # 公式理解
        # self.H2 += self.ANME.dot(self.gH2)
        # 消耗H2
        self.H2 += self.vSRB * self.SRB * -3.6  # 公式理解
        # self.H2 -= self.SRB.dot(self.cH2)
        # 物质扩散
        self.H2 = self.diffuse(self.H2)

    # 物质扩散通用方法
    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.1  # 扩散系数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.H2[x, y] ** 0.5 * self.CH4[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] + 4) / 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] + 4) / 8.314e-3 / 298)
                self.vSRB[x, y] = v0 * fk * ft
            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 += (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]
                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.GFE1[x, y] < 0:
            self.SRB += (self.vSRB - 0.11 / 365 / 24 / 60) * self.SRB  # 负增长是否对应负速率？吉布斯自由能是否还需要判断？
            # 分裂
            if self.SRB[x, y] > self.sSRB:
                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[x, y] *= 0.5  # TODO 分裂后生物量如何变？
                    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(100)
    # m.init_para2(8)

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