#coding: UTF-8

# Code for MSM-kf and MSM-dg imprement

import random
import math
import numpy as np
import gc

CKD_MORTARILTY = [1, 1, 1.2, 1.8, 3.2, 5.9] # CKD1, 2, 3a, 3b, 4, 5(predialysis)
CKD_QOL = [0, 0, 0, 0, 0, 0, 0] # CKD1 ,2 ,3a ,3b ,4 ,5 ,5D
CKD_COST = [0, 0, 0, 0, 0, 0, 0] # CKD1 ,2 ,3a ,3b ,4 ,5 ,5D
SIM_INT = 4 # Simulation interval (times per year)

class Person_basic:
        """basic person"""
        def __init__(self):
            self.age = 20
            self.sex = 0 # 0:male, 1:female
            self.lifestatus = 1 # 0:dead, 1:live
            self.qol = 0
            self.cost = 0
            self.simint = SIM_INT

        def ageset(self, x):
            self.age = x

        def sexset(self, x):
            self.sex = x

        def ageplus(self):
            self.age = self.age + (1 / self.simint) * self.lifestatus

        def addqol(self, x):
            self.qol = self.qol + x

        def addcost(self, x):
            self.cost = self.cost + x

        def motRateBase(self): # Mortality rate of predialysis patients
            return (3.3381-self.sex * 1.7236) * \
            np.exp((0.0910+self.sex * 0.007) * self.age) / 100000 /4 #Ontario

            #return 4.2479 * np.exp(0.0916 * self.age) / 100000 /4 #Japan

        def motRateRRT(self): # Mortality rate of dialysis patients
            return 13 * np.exp(0.0599 * self.age) / 10000 /4

#-----------------------------------
class Person_FM(Person_basic): #inheriting Person_basic
    """Virtual patient of Fixed probability model"""
    def __init__(self):
        Person_basic.__init__(self)
        self.status = 1 # CKD grades
        self.transprob = [0,0,0,0,0,0,0] # Set default transition probabilities

    def settransprob(self, x1, x2, x3, x4, x5, x6):
        self.transprob[0] = x1 # 1 to 2
        self.transprob[1] = x2 # 2 to 3a
        self.transprob[2] = x3 # 3a to 3b
        self.transprob[3] = x4 # 3b to 4
        self.transprob[4] = x5 # 4 to 5
        self.transprob[5] = x6 # 5 to 5D

    def inputstatus(self,x):
        self.status = x

    def changestatus(self): # Personal simulation
        rnd = random.random() # Random number for life/die decision
        motb = self.motRateBase() # Mortality rate of predialysis patients

        # Life or die
        if self.lifestatus == 0: # Already dead
            pass
        elif self.status <= 6:
            if rnd < motb * CKD_MORTARILTY[self.status - 1]: # Predialysis
                self.lifestatus = 0
        elif self.status == 7: # Patients in dialysis
            if rnd < self.motRateRRT():
                self.lifestatus = 0
        else:
            raise NameError("Invalid Status Error")
        # Life or die ends

        if random.random() < self.transprob[self.status-1]: # State transition
            self.status = min(self.status + 1, 7)

        self.ageplus() # Age added
        self.addqol(CKD_QOL[self.status - 1]) # QOL added (optional)
        self.addcost(CKD_COST[self.status - 1]) # Cost added (optional)

    def export(self): # Output to personal time course matrix
        return [self.age, self.sex, self.status, self.lifestatus, self.qol, \
        self.cost]

#-----------------------------------
class Person_eDM(Person_basic): #inheriting Person_basic
    """Virtual patient of eGFR dependent model"""
    def __init__(self):
        Person_basic.__init__(self)
        self.eGFR = 90
        # Mean, Standard Deviation of eGFR decline (Annual)
        self.eGFRdec = [0.1,0.5]
        # eGFR decline speed of the patient, log-normal distribution assumed
        declogsigma = math.sqrt(math.log(self.eGFRdec[1] **2 / self.eGFRdec[0] **2 + 1))
        declogmu = math.log(self.eGFRdec[0]) - (declogsigma **2) / 2
        self.eGFRdecsp = random.lognormvariate(declogmu,declogsigma)

    def inputeGFR(self,x): #set initial eGFR value
        self.eGFR = x

    def seteGFRdec(self, x1,x2): #set mean/sd of eGFR decline speed
        self.eGFRdec[0] = x1
        self.eGFRdec[1] = x2
        declogsigma = math.sqrt(math.log(self.eGFRdec[1] **2 / self.eGFRdec[0] **2 + 1))
        declogmu = math.log(self.eGFRdec[0]) - (declogsigma **2) / 2
        self.eGFRdecsp = random.lognormvariate(declogmu,declogsigma)

    def changestatus(self):
        rnd = random.random()
        motb = self.motRateBase()

        # Life or die
        if self.lifestatus == 0: # Already dead
            pass
        elif self.eGFR >= 7: # Predialysis patients
            if rnd < motb * CKD_MORTARILTY[ckdconv(self.eGFR) - 1]:
                self.lifestatus = 0
        elif self.eGFR >= 0: # Dialysis patients
            if rnd < self.motRateRRT():
                self.lifestatus = 0
        else:
            raise NameError("Invalid Status Error")
        # Life or die ends

        self.ageplus() # Age added
        self.addqol(CKD_QOL[ckdconv(self.eGFR) - 1]) # QOL added (optional)
        self.addcost(CKD_COST[ckdconv(self.eGFR) - 1]) # Cost added (optional)
        self.eGFR = max(self.eGFR - self.eGFRdecsp / self.simint, 0)

    def export(self):
        return [self.age, self.sex, self.eGFR, self.lifestatus, self.qol, \
        self.cost]


#-----------------------------------
def Fixed_Model_sim(coh, years, x1, x2, x3, x4, x5, x6,\
initAge = 40, initStatus = 2, female = 0.5):
# coh = virtual cohort size, years = simulation length

    psn = [] # List of virtual patients

    for i in range(0, coh):
        psn.append(Person_FM())

    list_psn = [] # Virtual cohort time course matrix

    for i in range(0, math.ceil(coh * female)): # Female patients
        psn[i].sexset(1)

    for i in range(0, coh):
        list_t = [] # Personal time course matrix
        psn[i].ageset(initAge) # Set initial age
        psn[i].inputstatus(initStatus) # Set initial CKD grade
        psn[i].settransprob(x1, x2, x3, x4, x5, x6) # Set transition probabilities

        for j in range(0, years * SIM_INT): # Simulation until time expires
            list_t.append(psn[i].export())
            psn[i].changestatus()

        list_psn.append(list_t)

    del psn # Garvage collection for avoiding memory leak
    gc.collect()

    return list_psn
#-----------------------------------
def eGFR_dependent_sim(coh, years, initAge = 40, initeGFR = 70, female = 0.5, \
decMean = 2.7,decStd = 2.8):
    psn = [] # List of virtual patients

    for i in range(0, coh):
        psn.append(Person_eDM())

    list_psn = [] # Vertual cohort time course matrix

    for i in range(0,math.ceil(coh * female)): # Female patients
        psn[i].sexset(1)

    for i in range(0, coh):
        list_t = [] # Personal time course matrix
        psn[i].ageset(initAge) # Set initial age
        psn[i].inputeGFR(initeGFR) # Set initial eGFR
        psn[i].seteGFRdec(decMean, decStd) # Set eGFR decline parameters

        for j in range(0, years * SIM_INT):
            list_t.append(psn[i].export())
            psn[i].changestatus()
        list_psn.append(list_t)

    del psn # Garvage collection for avoiding memory leak
    gc.collect()

    return list_psn

#-----------------------------------
def ckdconv(x): #conversion from eGFR to CKD grades
    if x >= 90: #G1
        return 1
    elif x >= 60: #G2
        return 2
    elif x >=45: #G3a
        return 3
    elif x >= 30: #G3b
        return 4
    elif x >= 15: #G4
        return 5
    elif x >= 7: #G5
        return 6
    elif x >= 0: #G5D
        return 7
    else:
        raise NameError("CKD Conv Error")
#-----------------------------------
# conversion from eGFR dependent model to transition probabilites
def eGFR2tp(nsize, splesize, x1, x2, years = 60, initAge = 40, \
initeGFR = 70, female = 0.5):
# x1: mean, x2: SD of eGFR decline speed
# splesize: sample size for transition probabilites calculation

    # First comes eGFR dependent model
    msim = eGFR_dependent_sim(coh = nsize, years = years, \
    decMean = x1, decStd = x2, \
    initAge = initAge, initeGFR = initeGFR, female = female)

    res = []

    for i in range(0, splesize): # Random sampling
        a = random.random() # Patient randomly selected
        a = math.ceil(a*nsize-1)
        tim = random.random() # Time randomly selected
        tim = math.ceil(tim*years-1)

        while msim[a][tim][3] == 0: # If already dead, re-try random sampling
            a = random.random()
            a = math.ceil(a*nsize-1)
            tim = random.random()
            tim = math.ceil(tim*years-1)

        tmp = [ckdconv(msim[a][tim][2]), ckdconv(msim[a][tim+1][2]), \
        msim[a][tim+1][3]] # Stage, Stage of next period, life of next period
        res.append(tmp)

    p = []

    for i in range(1,8):

        if res.count([i,i,1]) + res.count([i,i+1,1]) + res.count([i,i,0]) + \
        res.count([i,i+1,0]) == 0:
            tp = 0 # When particular stage patient does not exist
        else:
            tp = res.count([i,i+1,1]) / \
            (res.count([i,i,1]) + res.count([i,i+1,1]) \
            + res.count([i,i,0]) + res.count([i,i+1,0]))

        p.append(tp) # Transition probabilities

    return (p, msim)

#-----------------------------------
# For Maintenance
if __name__ == "__main__":
    aa,b = eGFR2tp(1000,1000, 4.8, 7.5, 60, 40, 70, 0.5)
    print(aa)