import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import shapiro
from scipy.stats import chisquare
from scipy.stats import t
from scipy.stats import norm
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import fileIO

DEFAULT_NUM = 10000
EPSILON = 1e-20
INFINITY = 1000000
NEG_INFINITY = -1000000
def almostEqual(a,b,epsilon = EPSILON):
    return np.abs(a-b) < epsilon

def makeDivisible(a, d = 2):
    if a % d != 0:
        return d - a % d + a
    else:
        return a

def findCommonInterval(list1, list2, as_series = True):
    """
    find the common interval of list  1 and list 2

    ----
    Parameters:
    ----

    as_series: bool
    boolean indicating whether list1 and list2 are arithmetic sequence, if so, more efficient method can be used

    return: low1,up1, low2,up2, s.t list1[low1:up1] ~ list2[low2:up2]

    ------
    To do:
    ------
    create version for non-as series, and check also whether the given series are arithmetic
    check also that the two lists can have an intersection, also the calculation for high_small and low_large is very sloppy
    """
    # find which one is larger
    if (list1[1] < list2[1]):
        larger = list2
        smaller = list1
        inOrder = True
    else:
        larger = list1
        smaller = list2
        inOrder = False
    # epsilon = INFINITY
    # for count in range(len(list1)):
    #     difference = abs(list(1) - )

    ## alternative
    ## assuming arithmetic_sequence
    if (as_series):
        dstep_small = smaller[1] - smaller[0]
        dstep_large = larger[1] - larger[0]
        low_small = int(dstep_large / dstep_small) - 1
        high_small = len(smaller)  ### sloppy
        low_large = 1
        high_large = int(smaller[-1] / dstep_large) + 1
        if (inOrder):
            return low_small, high_small, low_large, high_large
        else:
            return low_large, high_large, low_small, high_small
    else:

        raise Exception("Non arithmetic increase not supported yet!")

def findCommonPoints(list1, list2, as_series = True):
    """
    find the common points (in the sense of almostEqual) of two lists

    -----
    Parameters:
    -----

    return: index1List, index2List,

    ------
    To do:
    ---------
    The case for non as_series
    """
    # find which one is larger
    if (list1[1] < list2[1]):
        larger = list2
        smaller = list1
        inOrder = True
    else:
        larger = list1
        smaller = list2
        inOrder = False
    # epsilon = INFINITY
    # for count in range(len(list1)):
    #     difference = abs(list(1) - )

    ## alternative
    ## assuming arithmetic_sequence
    if (as_series):
        dstep_small = smaller[1] - smaller[0]
        dstep_large = larger[1] - larger[0]
        factor = int(dstep_large / dstep_small)

        low_large = 0
        low_small = int((larger[low_large] - smaller[0]) / dstep_small)
        high_small = len(smaller)  ### sloppy

        smallIndex = list(range(low_small, high_small, factor))

        high_large = int((smaller[-1] - larger[low_large]) / dstep_large) + 1
        largeIndex = list(range(low_large, high_large))
        if (inOrder):
            return smallIndex, largeIndex
        else:
            return largeIndex, smallIndex
    else:

        raise Exception("Non arithmetic increase not supported yet!")

def merge2TimeScale(df1, df2, columns = ["MSD_xy_mean","MSD_z"]):
    """
    given two dataframe saving the msd for two different time scales, merge the two together

    ------
    To do:
    ------
    Check first that df1 and df2 are of the same composition and molecule!

    ------
    Parameters:
    ----------

    df1: dataframe,

    df2: dataframe

    columns: list like, saving the columns to be merged

    return : dataframe with columns as specified in the variable columns

    ------
    Assuming
    -------
    dataframe has columns called "time"

    """

    ## To do, check if it is reasonable to merge these two

    ## find common interval
    l1, l2 = findCommonPoints(df1["time"], df2["time"], as_series = True)
    df1 = df1.iloc[l1]
    df2 = df2.iloc[l2]
    if (len(df1) != len(df2)):
        print(l1)
        print(l2)
        raise Exception("findCommonPoints error, findCommonPoints returned list1 and 2 not equal length!")
    for col in columns:
        minDiff = INFINITY
        minIndex = -1
        for i in range(len(df1)):
            difference = abs(df1.iloc[i][col] - df2.iloc[i][col])
            if (difference < minDiff):
                minDiff = difference
                minIndex = i
        timePt = df1.iloc[minIndex]["time"]

def chooseIndex(indices, num, lowPercentile = 0, upPercentile = 1):
    """
    choose indices from [lowPercentile * len, upPercentile * len]:

    ------
    Parameters:
    -------

    indices: array like

    num: int

    lowPercentile: float

    upPercentile: float

    """

    lenIndices = len(indices)
    lowerBound = int(np.floor(lenIndices * lowPercentile))
    upperBound = int(np.floor(lenIndices * upPercentile))
    # upperBound will not be reached so no need to add 1,
    # think for example upPercentile = 1, then upPercentile == lenIndices, which should not be reached as an index of the list
    # but lowerBound = 0 could be reached.
    totalLen = upperBound - lowerBound
    if (totalLen < 0):
        raise Exception("totalLen < 0!")
    if ( num >= totalLen):
        return list(range(lowerBound, upperBound))
    else:
        indexPrim = np.random.choice(range(lowerBound, upperBound), num , replace= False)
        indexPrim.sort()
        return indexPrim

def rvNormalityTest(rvSeries, alpha = "0.05"):
    """
    Test if a given series of random variables is normal distributed based on shapiro wilk

    ------
    Parameter
    -------

    rvSeries:
    array like structure, series of random variables

    alpha: float
    the alpha value for the hypothesis test

    """
    stat, p = shapiro(rvSeries)
    if p > alpha:
        return True
    else:
        return False

def rvTest(rvSeries, alpha = "0.05", distribution = "normal"):
    """
    Test if a given series of random variables is distributed based according to the scheme specified distribution

    ------
    Parameter
    -------

    rvSeries:
    array like structure, series of random variables

    alpha: float
    the alpha value for the hypothesis test

    """

    if (distribution == "normal"):
        stat, p = shapiro(rvSeries)
    elif (distribution == "chisquare"):
        stat, p = chisquare(rvSeries)

    if p > alpha:
        return True
    else:
        return False

def dfNormalityTest(df, num = DEFAULT_NUM, lowPercentile = 0.02, upPercentile = 0.1, alpha = 0.05, normalityTestCol = ["MSD_xy"]):
    """
    Test if a dataframe from the msd analysis, in particular the "MSD_xy"

    -------
    Parameter
    -------

    df: pandas.DataFrame

    num: int

    alpha: float

    """

    ## select only the MSD_xy and time
    ### maybe could insert at the beginning

    ### To do
    ### add concentration, molecule, info
    timeStep = df["samplingrate"].unique()
    dataFrameCol = ["concentration", "samplingrate", "molecule", "lowPercentile","upPercentile", "alpha"] + normalityTestCol
    testResultFinalDf = pd.DataFrame(columns = dataFrameCol)
    concentrationLists = df["concentration"].unique()
    for samplingRate in timeStep:
        df_samplingLevel = df[df["samplingrate"] == samplingRate]
        for concentration in concentrationLists:
            moleculeConcentrationDict = fileIO.getMoleculeConentration(concentration)
            columnConcentration = ["time"] + normalityTestCol + ["concentration", "molecule"]
            df_concentrationLevel = df_samplingLevel[columnConcentration]
            df_concentrationLevel = df_concentrationLevel[df_concentrationLevel["concentration"] == concentration]
            for molecule in moleculeConcentrationDict.keys():
                df_moleculeLevel = df_concentrationLevel[df_concentrationLevel["molecule"] == molecule]
                testResultDict = {}
                # initialize testResultDict
                for elem in normalityTestCol:
                    testResultDict[elem] = 0

                columns = ["time"] + normalityTestCol
                df_prune = df_moleculeLevel[columns]
                timeLag = df_prune["time"].unique()
                Indices = chooseIndex(list(range(len(timeLag))), num, lowPercentile = lowPercentile, upPercentile= upPercentile)
                # perhaps there should be more molecules, but this dataframe only include one
                # like POPC75_CHOL25, and only CHOL is present
                if (len(Indices) < 1):
                    continue
                for index in Indices:
                    tempDF = df_prune[almostEqual(df_prune["time"], timeLag[index])]
                    for elem in normalityTestCol:
                        rvSeries = tempDF[elem]
                        isNormal = rvNormalityTest(rvSeries, alpha)
                        if (isNormal):
                            testResultDict[elem] += 1
                # print("the testResultDict absolute: ")
                # print(testResultDict)
                totalNum = len(Indices)
                for key in testResultDict.keys():
                    testResultDict[key] /= totalNum
                # print("------------------------------")
                # print("after modification:")
                # print("the testResultDict absolute: ")
                # print(testResultDict)
                # print("-----------------------------")
                ## print result:
                print ("running normality test with time range :")
                print("({0:.4f}, {1:.4f})".format(timeLag[Indices[0]], timeLag[Indices[-1]]))
                print("total number of samplings: ")
                print(totalNum)
                print("alpha: " + str('{:.5f}'.format(alpha)))
                print("rate of passing:")
                # ["concentration", "samplingrate", "molecule", "lowPercentile","upPercentile"] + normalityTestCol
                df_testResult = pd.DataFrame({
                    "concentration": [concentration],
                    "samplingrate": [samplingRate],
                    "molecule": [molecule],
                    "lowPercentile": [lowPercentile],
                    "upPercentile":[upPercentile],
                    "alpha": [alpha]
                }
                )
                for key, value in testResultDict.items():
                    print("%-10s %10.5f"%(key, value))
                    df_testResult[key] = value
                testResultFinalDf = testResultFinalDf.append(df_testResult)
    return testResultFinalDf
def dfTest(df, num = DEFAULT_NUM, lowPercentile = 0.02, upPercentile = 0.1, alpha = 0.05, normalityTestCol = ["MSD_xy"], distribution = "normal"):
    """
    Test if a dataframe from the msd analysis, in particular the "MSD_xy"

    -------
    Parameter
    -------

    df: pandas.DataFrame

    num: int

    alpha: float

    """

    ## select only the MSD_xy and time
    ### maybe could insert at the beginning

    ### To do
    ### add concentration, molecule, info
    timeStep = df["samplingrate"].unique()
    dataFrameCol = ["concentration", "samplingrate", "molecule", "lowPercentile","upPercentile","alpha", "distribution"] + normalityTestCol
    testResultFinalDf = pd.DataFrame(columns = dataFrameCol)
    concentrationLists = df["concentration"].unique()
    for samplingRate in timeStep:
        df_samplingLevel = df[df["samplingrate"] == samplingRate]
        "df_samp"
        for concentration in concentrationLists:
            moleculeConcentrationDict = fileIO.getMoleculeConentration(concentration)
            columnConcentration = ["time"] + normalityTestCol + ["concentration", "molecule"]
            df_concentrationLevel = df_samplingLevel[columnConcentration]
            df_concentrationLevel = df_concentrationLevel[df_concentrationLevel["concentration"] == concentration]
            for molecule in moleculeConcentrationDict.keys():
                df_moleculeLevel = df_concentrationLevel[df_concentrationLevel["molecule"] == molecule]
                testResultDict = {}
                # initialize testResultDict
                for elem in normalityTestCol:
                    testResultDict[elem] = 0

                columns = ["time"] + normalityTestCol
                df_prune = df_moleculeLevel[columns]
                timeLag = df_prune["time"].unique()
                Indices = chooseIndex(list(range(len(timeLag))), num, lowPercentile = lowPercentile, upPercentile= upPercentile)
                # perhaps there should be more molecules, but this dataframe only include one
                # like POPC75_CHOL25, and only CHOL is present
                if (len(Indices) < 1):
                    continue
                for index in Indices:
                    tempDF = df_prune[almostEqual(df_prune["time"], timeLag[index])]
                    for elem in normalityTestCol:
                        rvSeries = tempDF[elem]
                        distriTrue = rvTest(rvSeries, alpha, distribution= distribution)
                        if (distriTrue):
                            testResultDict[elem] += 1
                # print("the testResultDict absolute: ")
                # print(testResultDict)
                totalNum = len(Indices)
                for key in testResultDict.keys():
                    testResultDict[key] /= totalNum
                # print("------------------------------")
                # print("after modification:")
                # print("the testResultDict absolute: ")
                # print(testResultDict)
                # print("-----------------------------")
                ## print result:
                print ("running test with time range :")
                print("({0:.4f}, {1:.4f})".format(timeLag[Indices[0]], timeLag[Indices[-1]]))
                print("total number of samplings: ")
                print(totalNum)
                print("alpha: " + str('{:.5f}'.format(alpha)))
                print("rate of passing:")
                # ["concentration", "samplingrate", "molecule", "lowPercentile","upPercentile"] + normalityTestCol
                df_testResult = pd.DataFrame({
                    "concentration": [concentration],
                    "samplingrate": [samplingRate],
                    "molecule": [molecule],
                    "lowPercentile": [lowPercentile],
                    "upPercentile":[upPercentile],
                    "alpha": [alpha],
                    "distribution": [distribution]
                }
                )
                for key, value in testResultDict.items():
                    print("%-10s %10.5f"%(key, value))
                    df_testResult[key] = value
                testResultFinalDf = testResultFinalDf.append(df_testResult)
    return testResultFinalDf

def histPlot(df, timePoints, timeNum , particleNum, bins = 50, normalityTestCol = ["MSD_xy"]):
    """
    plot the histogram of df

    -------
    Parameters:
    -------

    timePoints: array-like structure
    the time points from which to choose points to plot

    timeNum: int
    number of time points to choose

    particleNum: int
    number of particles to choose from

    """

    ## make sure the timePoints are unique
    timePoints = list(set(timePoints))
    concentrationLists = df["concentration"].unique()

    # select certain particles
    # calculate the total timeNum of trajectories
    numTrajectories = len(df["trajectory"].unique())
    numTrajectories = min(numTrajectories, particleNum)
    trajectoryIndex = chooseIndex(df["trajectory"],particleNum)
    trajectories = np.array(df["trajectory"])[trajectoryIndex]
    # subset the df
    df = df[df["trajectory"].isin(trajectories)]

    columns = ["time"] + normalityTestCol
    df_prune = df[columns]
    timePointsNum = len(timePoints)

    timeNum = makeDivisible(timeNum)
    numberOfPoints = min(timePointsNum, timeNum)
    if (numberOfPoints > timePointsNum):
        timeLists = np.random.choice(timePoints, size = numberOfPoints, replace = False)
    else:
        timeLists = timePoints

    timeLists.sort()
    # df_prune = df_prune[df_prune[almostEqual(df_prune["time"], timeLists)]]
    ## to do
    ## adapt plotting
    ncol = 2
    nrow = int(timeNum / ncol)

    fig,axes = plt.subplots(nrow,ncol)
    fig.tight_layout(pad = 5.0)
    fig.suptitle("total number of particles chosen: " + str(numTrajectories), fontsize = 20)
    count = 1
    for elem in normalityTestCol:
        for time in timeLists:
            rvVariables = df_prune[almostEqual(df_prune["time"], time)][elem]
            plt.subplot(nrow, ncol, count)
            rvVariables.plot.hist(bins = bins ) #ax = axes[count, 0]
            plt.title("time:{0:.2f}".format(time))
            plt.xlabel(elem)
            ## figSub.autofmt_xdate()
            count += 1

    plt.show()

def testHist(df, timeNum = 8, particleAverageNum = 8, normalityTestCol = ["MSD_xy"]):
    """

    given a dataframe, genearte a historam visualisation of the dataframe

    """

def diffArr(arr, method = "forward"):
    """

    calculates the difference between each elements of arr

    -------
    Parameters:
    ------

    method: str
    method used for diffArr, "forward", "central", "backward"

    -------
    To do:
    -------
    for central diff, how to treat the first element, here use upward diff
    """
# not sure if some arr starts indexing with 0
# for example, doing a slicing of a dataframe, will cause arr starts at index 1
# df_prune = df[1:]
# diffArr(df_prune["MSD_xy"])  <---- index starts at 1
    lenArr = len(arr)
    arTemp = [None] * (lenArr - 1)
    count = 0
    if (method == "forward"):
        for val in arr:
            if (count == 0):
                oldval = val
            else:
                arTemp[count - 1] = val - oldval
                oldval = val
            count += 1
    elif (method == "central"):
        for val in arr:
            if (count == 0):
                oldval = val
            elif (count == 1):
                arTemp[0] = val - oldval # the diff for the first element use forward
            elif(count >= 2):
                arTemp[count - 1] = val - oldval
                oldval = arr[count - 1]
            count += 1
    elif (method == "backward"):
        raise Exception("temporarily not supported for backward")
    else:
        raise Exception("unknown type of method for calculation diff: " + method)

    return np.array(arTemp)

def getSlope(xArr, yArr, method = "forward"):
    """
    get the slope of yArr w.r.t xArr,

    ------
    Parameters:
    ------

    method: str
    method use for calculating slope, "forward", "central", "backward"

    ---
    To do:
    ----
    1.  realize the "backward" scheme
    2.  does central makes sense, if x is not arithmetic sequence?
    3.  could relax np.array(diffArr(...)) because diffArrr already ensures the returned is ndarray
    """
    if (len(xArr) != len(yArr)) :
        raise Exception("x and y not equal length, cant calculate the slope")
    if (method == "forward"):
        xArr_diff = np.array(diffArr(xArr, method = "forward"))
        yArr_diff = np.array(diffArr(yArr, method = "forward"))
    elif (method == "central"):
        xArr_diff = np.array(diffArr(xArr, method = "central"))
        yArr_diff = np.array(diffArr(yArr, method = "central"))
    elif (method == "backward"):
        xArr_diff = np.array(diffArr(xArr, method = "backward"))
        yArr_diff = np.array(diffArr(yArr, method = "backward"))
    else:
        raise Exception("Unknown type of method for slope calculation: " + method)

    slopes = yArr_diff/xArr_diff
    return slopes
def eliminateZeroTime(df):
    """
    eliminate the 0 time row, (doing log will cause problem)

    -----
    To do
    ------
    figure out how to drop 0 that emerge at rows other than the first line
    """

    if (almostEqual(df.iloc[0,:]["time"], 0)):
        return df[1:]
    else:
        return df

def findLowUpIndex(arr, lowBound = NEG_INFINITY, upBound = INFINITY, startIndex = -1, endIndex = -1, userIsKing = True):
    """
    find the left, right index of an array, that has a continuous range lying between [lowBound, upBound]
    only count the first time this range is hit, and with consideration to the startIndex

    -------
    Parameters:
    -------
    arr: array like structure

    lowBound: float

    upBound: float

    startIndex: int

    endIndex: int, -1 default meaning no constraint

    userIsKing: boolean, default True :-)
    when there is collision of the lowBound and the available min, if False, use the available min, else the considering the user given and choose an optimal

    ------
    To do:
    ------
    optimise the code for lowBound == NEG_INFINITY, upBOUND == INFINITY
    """

    # make sure it is a numpy.array
    arr = np.array(arr)
    boolArr = (arr < upBound) & (arr > lowBound)
    leftIndex = -1
    rightIndex = -1
    leftFound = False
    rightFound = False
    continuous = True

    if (endIndex > 0):
        endIndex = min(endIndex, len(arr))
    else:
        endIndex = len(arr)

    if (startIndex < 0):
        startIndex = 0
    arrSlice = arr[startIndex : endIndex]
    arrSlice_low = min(arrSlice)
    arrSlice_max = max(arrSlice)

    if ( not userIsKing):
        lowBound = arrSlice_low
        upBound = arrSlice_max
    else:
        if (lowBound < arrSlice_max ):
            lowBound = max(lowBound, arrSlice_low)
        if (upBound > arrSlice_low):
            upBound = min(upBound, arrSlice_max)
    for i in range(len(boolArr)):
        if (i >= startIndex):
            if ((boolArr[i] == True) and (leftFound == False)):
                leftFound = True
                leftIndex = i
            elif (leftFound == True):
                # need to first check if the leftIndex is already found
                if((boolArr[i]) and (not rightFound) and (continuous)):
                    rightIndex = i
                elif (rightFound == False):
                    # boolArr[i] == False
                    rightFound = True
                    continuous = False
                    break
        if (i > endIndex):
            if (leftFound == False):
                raise Exception("could not find any index, leftIndex already beyond the array")
            else:
                break
    print("lowBound: " + str(lowBound))
    print("upBound: " + str(upBound))
    print("leftIndex: " + str(leftIndex))
    print("rightIndex: " + str(rightIndex))
    return leftIndex, rightIndex

def getLogLogSlope(df, yColumn = ["MSD_xy_mean"], xColumn = ["time"], method = "forward"):
    """
    find the slope k which is given by log(y) = k * log(x) + b

    ------
    Parameters:
    ------

    method: str
    method used for the calculation of slopes, "forward", "central", "backward"
    """
    ## eliminate possible zero rows

    if (len(yColumn) > 1 or len(xColumn) > 1):
        raise Exception("Only support for one dimensional array!")
    if (len(yColumn) < 1 or len(xColumn) < 1):
        raise Exception("no entry in yColumn or xColumn!")
    df_prune = eliminateZeroTime(df)

    for xCol in xColumn:
        for yCol in yColumn:
            tLog = np.log(df_prune[xCol])
            msdLog = np.log(df_prune[yCol])
            tLogDiff = diffArr(tLog, method = method)
            msdLogDiff = diffArr(msdLog, method = method)
            slopes = msdLogDiff / tLogDiff
            return slopes

def plotLogLog(df, yColumn = ["MSD_xy_mean"], label = ""):
    """
        print the loglog
    """
    df_prune = eliminateZeroTime(df)
    tLog = np.log(df["time"])
    for col in yColumn:
        yLog = np.log(df[col])
        plt.loglog(np.exp(tLog),np.exp(yLog), label = label)

def divide3Region(df, columns = ["MSD_xy_mean"],threshold2_low = 0.3,threshold2_high = 0.8, \
    threshold1_low = 1.6, threshold3_low = 0.8, region2start = -1, region3start = -1, secondDerivThreshold = 0.6, \
        indexJump = 20, method = "forward"):
    """

    very specific for MSD analysis, given the msd, divide the series into three regions
    assume df has "time" as one of its column name
    ------
    Parameters
    -------

    secondDerivThreshold: float
    set the threshold for the secondDerivative of the slopes

    return: three lists and the slopes
    example usage l1,l2,l3,slopes = divide3Region(df, columns = ["MSD_xy"])
    df_1 = df.iloc[l1]

    ------
    To do
    ------

    1.  make several columns possible, make the search adaptive, do not require user to guess which threshold to set, maybe could try to use second order difference
    2.  region3Start, need  to check not too large?

    """
    # threshold1_low = 1.6
    # threshold3_low = 0.8
    if (threshold3_low < threshold2_high + EPSILON):
        threshold3_low += 0.09
    threshold3_high = 1.2
    df_prune = eliminateZeroTime(df)
    if (len(columns) > 1):
        raise Exception("currenlty only division works for one column")

    for col in columns:
        # timeSeries = df_prune["time",col]
        # tLog = np.log(df_prune["time"])
        # msdLog = np.log(df_prune[col])
        # tLogDiff = diffArr(tLog)
        # msdLogDiff = diffArr(msdLog)
        # slopes = msdLogDiff / tLogDiff
        slopes = getLogLogSlope(df_prune, yColumn = [col], xColumn= ["time"], method = method)
        slopesDiff = diffArr(slopes)
        tLog = diffArr(df_prune["time"])[0:-1]
        slopesDeriv = slopesDiff / tLog
        print("slopesDeriv:")
        print(slopesDeriv)
        boolArr = np.abs(slopesDeriv) > secondDerivThreshold
        print("number > threshold : " + str(np.sum(boolArr)))
        print("indexArr:")
        # indexArr n * 1
        indexArr = np.where(boolArr)[0]
        print(indexArr)
        # find the region where the index jump occurs
        indexArrDiff = diffArr(indexArr)
        print("indexArrDiff:")
        print(indexArrDiff)
        boolDiffArr = np.abs(indexArrDiff) > indexJump
        index_indexArr = np.where(boolDiffArr)[0]
        print("index_indexArr:")
        print(index_indexArr)
        start2IndexCandidate = indexArr[index_indexArr[0]]
        print("startPoint for region2")
        print(start2IndexCandidate)

        start2Index = max(start2IndexCandidate, region2start)
        end2Index = indexArr[index_indexArr[0] + 1]

        ## find the first region
        l1Dummy, r1 = findLowUpIndex(slopes, lowBound= threshold1_low)
        l2,r2 = findLowUpIndex(slopes, lowBound=threshold2_low, upBound= threshold2_high, startIndex= start2Index, endIndex= end2Index)
        # have to make sure that l3 starts to count after the region 2

        region3start = max(region3start, r2)
        if (region3start >= len(df[col])):
            raise Exception("region3start too large, already exceeds the series length!")
        l3,r3Dummy = findLowUpIndex(slopes, lowBound= threshold3_low, upBound = threshold3_high, startIndex=region3start)
        list1 = list(range(0,r1 + 1))
        list2 = list(range(l2,r2 + 1))
        list3 = list(range(l3,len(slopes)))
        return list1, list2, list3, slopes

def simpleLinearRegression(xArr,yArr, confidence = 0.95):
    """
    do a linear regression s.t yArr_approx = k * xArr + b

    ----
    Parameters:
    ----

    return: dictionary which includes k,b,mean_squared_error,confidence level,  confidenceInterval of k, confidenceInterval of b
    """
    model = LinearRegression().fit(xArr,yArr)
    x = xArr.flatten()
    yPred = model.predict(xArr)
    n = len(x)
    if (n > 2):
        mean_sqre_error = mean_squared_error(yPred,yArr) * n / (n - 2)
    else :
        mean_sqre_error = mean_squared_error(yPred,yArr)
        print("less than 2 points, the estimator for variance sigma may be invalid")

    lxx = np.sum(np.square(x)) - np.square(np.mean(x)) * n
    if (lxx < 0):
        raise Exception("lxx < 0!")
    alpha = 1 - confidence
    alphaHalfQuantile = t(df = n).ppf(1 - alpha/2)
    k_interval_len = np.sqrt(mean_sqre_error/lxx)
    k = model.coef_[0]
    k_low = k - alphaHalfQuantile * k_interval_len
    k_high = k + alphaHalfQuantile * k_interval_len

    b = model.intercept_
    b_interval_len = np.sqrt(mean_sqre_error * (1/n + np.square(np.mean(x))/lxx))
    b_low = b - b_interval_len * alphaHalfQuantile
    b_high = b + b_interval_len * alphaHalfQuantile
    return {'k': k, 'b': b, 'mean_square_error': mean_sqre_error, 'confidence': confidence, 'k_interval': [k_low, k_high], 'b_interval': [b_low, b_high]}

def msdRegression(df, columns = ["MSD_xy_mean"], confidence = 0.95):
    """
    do a linear regression s.t log(columns) = k * log(time) + b

    ----
    Pay Attention:
    ----
    Assume df has no zero Time

    -----
    Parmaeters:
    -----

    return: dictionary which includes k,b,mean_squared_error,confidence level,  confidenceInterval of k, confidenceInterval of b

    ----
    To do:
    -----

    Making multiple columns possible

    """

    # pandas.Series do not support reshape function, and given a pandas.Series, the np.log() wont convert that to a ndarray
    # so need to convert to a ndarray first
    tLog = np.log(df["time"]).to_numpy().reshape((-1,1))

    for col in columns:
        msdLog = np.log(df[col])
        return simpleLinearRegression(tLog, msdLog, confidence)

def msdFittingPlot(df, columns = ["MSD_xy_mean"], threshold2_low = 0.3,threshold2_high = 0.8, \
    threshold1_low = 1.6, threshold3_low = 0.8, region2start = -1, region3start = -1, time_unit = "ps", msd_unit = "nm", outputUnit = "si", \
        potentKnown = True, confidence = 0.95, secondDerivThreshold = 0.6, method = "forward", title = "", directShow = True):
    """
    plot the msd and the curve fitting

    -----
    Parameters
    -----

    time_unit: str,
    indicate the units of the timescale, default ps

    ----
    To do:
    ----
    1. df should eliminate 0?
    2. making multiple columns possible
    3. annotate the curve with slope and intercept
    4. given units, figure out the constant to change to SI seconds
    """
    df = eliminateZeroTime(df)

    if (time_unit == "ps" and msd_unit == "nm"):
        # factor = 1e-12
        offset_alpha = 12 * np.log(10) # each intercept calculated need to add this offset * alpha, e.g MSD * 1e-18 = k (1e-12 * t)^a, logMSD = logk + a * log(t) - 12 a * log(10) + 18 log(10)
        offset_const = -18 * np.log(10)
    else:
        raise Exception("currently does not support other units")

    resultRegression = {}
    for col in columns:
        list1, list2, list3, slopes = divide3Region(df,columns = [col], threshold2_low = threshold2_low, \
            threshold2_high = threshold2_high, threshold1_low = threshold1_low, threshold3_low = threshold3_low, \
                region2start = region2start, region3start = region3start, secondDerivThreshold= secondDerivThreshold,\
                    method = method)
        indexLists = [list1,list2,list3]
        dfLists = [df.iloc[indexLists[var]] for var in range(3)]
        fig, axes = plt.subplots()
        # plot the whole msd
        legends = ["MSD","ballistic","subdiffusion","brownian"]
        tLog = np.log(df["time"])
        msdLog = np.log(df[col])
        plt.plot(np.exp(tLog),np.exp(msdLog), label = legends[0])
        # plt.plot(tLog,msdLog, label = legends[0])
        if (outputUnit == "si"):
            resultRegression['units'] = "si"
        else:
            resultRegression['units'] = "nm_ps"
        alpha_2Interval = []
        for i in range(3):
            regressionDict = msdRegression(dfLists[i], columns = [col], confidence = confidence)
            k = regressionDict['k']
            b = regressionDict['b']
            b_interval = np.array(regressionDict['b_interval'])
            b_siunit = b + offset_alpha * k + offset_const
            b_si_interval = b_interval + offset_alpha * k + offset_const
            ts = np.log(df["time"])
            ys = k * ts + b
            legend = legends[i + 1]
            if (i == 0):
                intercept_1 = np.exp(b_siunit)
                if (outputUnit == "si"):
                    b_1 = intercept_1
                else:
                    b_1 = np.exp(b)
                resultRegression['b1'] = b_1
                resultRegression['b1_interval'] = regressionDict['b_interval']
                k_1 = k
                resultRegression['k1'] = k_1
                resultRegression['k1_interval'] = regressionDict['k_interval']
            elif ( i == 1):
                intercept_2 = np.exp(b_siunit)
                alpha_2 = k
                alpha_2Interval = regressionDict['k_interval']
                resultRegression['alpha'] = k
                resultRegression['alpha_interval'] = alpha_2Interval
                legend += ":{0:.5f}".format(k)
            else:
                intercept_3 = np.exp(b_siunit)
                k_3 = k
                if (outputUnit == "si"):
                    print("for b_3 only nm_ps units")
                    # b_3 = np.exp(b)
                    if (potentKnown == False):
                        print("The result for the errorbar of 2dD may be inaccurate!")
                        b_3 = intercept_3
                        b_3_interval = b_si_interval
                    else:
                        b_3 = np.exp(b) * 1e-6
                        b_3_interval = b_interval # move the multiplication of 1e-6 at the end
                else:
                    b_3 = np.exp(b)
                    b_3_interval = b_interval
                resultRegression['b3'] = b_3
                resultRegression['b3_interval'] = np.exp(b_3_interval)
                if (outputUnit == "si" and potentKnown):
                    resultRegression['b3_interval'] *= 1e-6
                resultRegression['k3'] = k_3
                resultRegression['k3_interval'] = regressionDict['k_interval']
            # ts = np.log(dfLists[i]["time"])

            plt.loglog(np.exp(ts), np.exp(ys), label = legend, linestyle = "dashed")
            # plt.plot(ts, ys, label = legends[i + 1], linestyle = "dashed")
        # calculate diffusion * d
        if (outputUnit == "si"):
            si_factor = (1e-6)
            msd_factor = 1e-18
            t_factor = 1e-12
        else:
            si_factor = 1
            msd_factor = 1
            t_factor = 1
        D_from_raw_data = df.iloc[-101:-1][col]/(df.iloc[-101:-1]['time']) * si_factor
        variance_name = col.split("_")
        variance_name[-1] = "var"
        variance_name = "_".join(variance_name)
        # heuristic:
        D_varList = df.iloc[-101:-1][variance_name]
        Dcompact_var = np.sum(D_varList)/100 * msd_factor
        D_std = np.sqrt(Dcompact_var)/(2 * df.iloc[-50]['time'] * t_factor)
        sum =0
        for i in D_from_raw_data:
            sum = sum + i
        Dd = sum/100
        alpha = 1 - confidence
        alphaHalfQuantile = norm.ppf(1 - alpha/2)
        Dd_interval_len = D_std * alphaHalfQuantile
        resultRegression['Dd'] = Dd
        resultRegression['Dd_interval'] = [Dd - Dd_interval_len, Dd + Dd_interval_len]
        plt.legend()
        plt.title(title)
        if (directShow):
            plt.show()
    constant1 = df.iloc[0][col] * 1e6 / np.square(df.iloc[0]["time"])
    constant2 = df.iloc[1][col] * 1e6 / np.square(df.iloc[1]["time"])
    print("ballistic region, one point calculated const (d * kT/m):{0:.4e}, {1:.4e}".format(constant1, constant2))
    print("ballistic region, const (d * kT/m): {0:.4e}, potent: {1:.4f}".format(intercept_1, k_1))
    print("subdiffusion region, intercept: {0:.4e}, potent alpha: {1:.4f}".format(intercept_2, alpha_2))
    print("confidence: {0:.4f}".format(confidence))
    print("interval for alpha: ")
    print(alpha_2Interval)
    print("brownian region, const (2dD): {0:.4e}, potent: {1:.4f}".format(intercept_3, k_3))
    return resultRegression

if __name__ == "__main__":
    arr = np.array([2,3,4,5,6,8])
    arrT = diffArr(arr)
    print(arrT)

    charArr = ["hi","there","yo","yo1","haha"]
    tempIndx = chooseIndex(charArr,2, 0,0.5)
    print(tempIndx)

    ar1 = np.arange(4,20,4)
    ar2 = np.arange(1,20,1)

    l1, l2 = findCommonPoints(ar1,ar2)
    print(ar1)
    print(ar2)
    print(l1)
    print(l2)


size = float(input())

counter = 1
command = 0
while command != "End":
    command = input()
    if command == "End":
        print("Congratulations! All suitcases are loaded!")
        break
    box = float(command)

    if counter % 3 == 0:
        box = (box * 0.10) + box

    if box > size:
        print("No more space!")
        break
    size -= box
    counter += 1

print(f"Statistic: {counter - 1} suitcases loaded.")


f = []  # f[i][j]表示i到j之间最大加分
root = []   # root[i][j]表示i到j之间的根结点

def solve(n):
    for length in range(2, n+1):
        for i in range(1, n):
            j = i + length - 1
            if j <= n:
                temp = 0
                for k in range(i, j+1):
                    if temp < f[i][k-1] * f[k+1][j] + f[k][k]:
                        temp = f[i][k-1] * f[k+1][j] + f[k][k]
                        root[i][j] = k
                f[i][j] = temp
    print(f[1][n])

def pre_order(s, e, n):
    if s > e or s < 1 or e > n:
        return
    temp = root[s][e]
    print(root[s][e], end=' ')
    pre_order(s, temp-1, n)
    pre_order(temp+1, e, n)

if __name__ == '__main__':
    n = int(input())
    arr = [int(i) for i in input().split()]
    f = [[1 for i in range(n+2)]for j in range(n+2)]
    root = [[1 for i in range(n+2)]for j in range(n+2)]
    for i in range(1, n+1):
        f[i][i] = arr[i-1]
        root[i][i] = i
    solve(n)
    pre_order(1, n, n)


n = int(input())
arr = input()
arr = [int(x) for x in arr.split(" ")]
arr.sort()
arr.reverse()

odd = []
even = []
for i in arr:
    if i % 2 == 0:
        even.append(i)
    else:
        odd.append(i)

even1 = even.copy()
odd1 = odd.copy()

while True:
    if even1.__len__() > 0:
        del even1[0]
        if odd1.__len__() > 0:
            del odd1[0]
        else:
            break
    else:
        break
count1 = 0
for i in even1:
    count1 += i
for i in odd1:
    count1 += i
while True:
    if odd.__len__() > 0:
        del odd[0]
        if even.__len__() > 0:
            del even[0]
        else:
            break
    else:
        break
count2 = 0
for i in even:
    count2 += i
for i in odd:
    count2 += i
print(min(count1, count2))


from covid import Covid
import sys
import matplotlib.pyplot as plt

covid = Covid()
countries = covid.list_countries()
Country_name=input("ENTER THE COUNTRY NAME:  ").capitalize()

#get the country's virus data
virus_data=covid.get_status_by_country_name(Country_name)
print(virus_data)

#save data related to confirmed cases for the future reference
Active_cases=virus_data['confirmed']

#remove unnecesary values like latitude longitude,lastupdate
remove=['id','country','latitude','longitude','last_update','confirmed']

for i in remove:
    virus_data.pop(i)

#virusdata is dictionary data type
#we should split them as keys and values respectively

id=list(virus_data.keys())
value=[str(i) for i in virus_data.values()]

#piechart creation
plt.pie(value,labels=id,colors=['r','yellow','g'] ,autopct='%1.1f%%')
plt.title("COUNTRY : "+Country_name.upper()+"\n TOTAL ACTIVE CASES :" +str(Active_cases))
plt.legend()
plt.show()


import os
import sys
import csv
import time

try:
    import tweepy
except ImportError:
    print("[ERROR] Unable to import Tweepy module: can't run!")
    sys.exit()

class FollowerScraper:

    def __init__(self):
        self._status = ""
        self._followers = list()
        self._csv_tweet = ""
        self._csv_tweet_writer = ""
        self._curr_follower = ""

    def get_followers(self, api, target, data):
        try:
            if os.name == "nt":
                self._csv_tweet = open(
                    data, 'w', encoding='utf-8', newline='')
            else:
                self._csv_tweet = open(data, 'w')

            self._csv_tweet_writer = csv.writer(self._csv_tweet)
            self._csv_tweet_writer.writerow(
                ["username"])

        except Exception as e:
            print("[ERROR] Unable to prepare CSV files!")
            print("\n[Details]: ", e)
            sys.exit()

        while True:
            try:
                print("[*] Downloading '", target, "' followers list")
                print("[*] Please, wait...")
                for self._tweet in tweepy.Cursor(
                        api.followers,
                        screen_name=target,
                        count=200
                ).items():

                    self._curr_follower = self._tweet.screen_name
                    print(self._curr_follower)

                    try:
                        self._csv_tweet_writer.writerow([
                            self._curr_follower
                        ])

                    except Exception as e:
                        print(
                            "[ERROR] Unable to write tweets on file: ",
                            data, ", details: ", e)
                        sys.exit()

            except tweepy.error.RateLimitError as e:
                print("[ERROR: TWEEPY API] Too many requests. Wait some minutes.")
                time.sleep(15 * 60)
            except tweepy.TweepError as e:
                if e.api_code == 429:
                    print("[ERROR: TWEEPY API] Too many requests. Wait some minutes.")
                    time.sleep(15 * 60)
                else:
                    print("[ERROR: TWEEPY API] " + str(e.text))
                sys.exit()
            except Exception as e:
                print("[ERROR]: ", e)
                sys.exit()


import argparse
import os
import random
import time

import orjson
import gc

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

from PIL import Image, ImageDraw, ImageChops, ImageFilter
from enum import Enum
import sys
import numpy as np
from pathlib import Path
import torch
import glob
from dataclasses import dataclass, field, astuple
import string
from concurrent.futures import ThreadPoolExecutor
import facechecker
import webuiapi
import functools
import SSIM_PIL
from moviepy.video.io.ImageSequenceClip import ImageSequenceClip

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # yolov5 strongsort root directory
WEIGHTS = ROOT / 'weights'

if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
if str(ROOT / 'yolov5') not in sys.path:
    sys.path.append(str(ROOT / 'yolov5'))  # add yolov5 ROOT to PATH
if str(ROOT / 'trackers' / 'strong_sort') not in sys.path:
    sys.path.append(str(ROOT / 'trackers' / 'strong_sort'))  # add strong_sort ROOT to PATH

if str(ROOT / 'rife') not in sys.path:
    sys.path.append(str(ROOT / 'trackers' / 'strong_sort'))

if r"TikTokDiffusion\rife" not in sys.path:
    sys.path.append(r"TikTokDiffusion\rife")

if r"rife" not in sys.path:
    sys.path.append(r"rife")

ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

import RIFE
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.dataloaders import LoadImages
from yolov5.utils.general import (LOGGER, Profile, check_img_size, non_max_suppression, scale_boxes, check_requirements,
                                   print_args, scale_segments)
from yolov5.utils.torch_utils import select_device
from utils.segment.general import masks2segments, process_mask, process_mask_native
from trackers.multi_tracker_zoo import create_tracker

@torch.no_grad()
def run(
        output_path,
        default_save_folder,
        fps,
        sampler,
        frame_gen_all,
        frame_gen_auto,
        cfg,
        height,
        width,
        seed,
        denoising,
        pipe,
        frame_gen_even,
        restore_faces,
        server,
        audio_source,
        control_net_model,
        control_net_lowvram,
        control_net_processor,
        source: Path,
        yolo_weights=WEIGHTS / 'yolov5x-seg.pt',  # model.pt path(s),
        reid_weights=WEIGHTS / 'osnet_x0_25_msmt17.pt',  # model.pt path,
        tracking_method='strongsort',
        imgsz=(640, 640),  # inference size (height, width)
        conf_thres=0.5,  # confidence threshold
        iou_thres=0.5,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        save_crop=False,  # save cropped prediction boxes
        classes=0,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # augmented inference
        visualize=False,  # visualize features
        half=False,  # use FP16 half-precision inference
        dnn=False,  # use OpenCV DNN for ONNX inference
        retina_masks=True,
        debug=False,
        skip_if_exists="",
        check_faces=False,
        prompt_by_body=False,
        steps=20,
        save_format=".png",
        stable_diffusion_folder="",
        face_match_sensitivity=1,
        func_prompt_list=None,
        ignore_if_no_prompt=True,
        fallback=None,
        loopback=1,
        inject_last_image=dict,
        variation_scaling=0.0,
        seed_scaling=False,
        ssim_retry=1,
        ssim_threshold=0.87,
        max_person_size=150,
        skip_ratio=.80,
        max_frame_skip=2,
        use_control_net=False,
        negative_prompt="",
):

    print(inject_last_image)
    org_max_person_size = max_person_size
    max_person_size = org_max_person_size*org_max_person_size
    source = str(source)
    source_folder = os.path.dirname(source)
    dataset = []
    name_prompt = func_prompt_list[0] if len(func_prompt_list[0]) < 20 else func_prompt_list[0][:20]+"..."
    id_set = set()

    if check_faces: #this makes the prompts work from right to left which is how prompt to body works and is more intuitive
        func_prompt_list = list(reversed(func_prompt_list))

    if skip_if_exists:
        save_folder = os.path.join(output_path, skip_if_exists)
    else:
        save_folder = default_save_folder

    class SaveType(Enum):
        DEBUG_IMAGE = "debug"
        UNFOUND_IMAGE = "unfound"
        FOUND_IMAGE = "found"
        FACES = "faces"

    class PersonJob:
        class_save_format = save_format
        class_full_save_folder_path = stable_diffusion_folder

        def __init__(self, org_image_path):
            self.org_image_path = org_image_path
            self.org_image_array = None
            self.segment = None

            filename = os.path.splitext(os.path.basename(self.org_image_path))[0]
            safe_filename = name_cleaner(filename)

            self.save_name_no_ext = safe_filename
            self.full_save_path_with_ext = os.path.join(save_folder, self.save_name_no_ext+f"{save_format}")
            self.mask_with_alpha = None
            self.masked_image = None
            self.prompt = None
            self.generated_times = 0
            self.found_face = False
            self.label = None
            self.face_index = None
            self.face = None
            self.track_label = None

        def save_debug_image(self, folder: SaveType, count=0):
            debug_folder = os.path.join(save_folder, "debugging", folder.value)

            try:
                os.makedirs(debug_folder)
            except FileExistsError:
                pass

            save_full_file_name = os.path.join(debug_folder, self.save_name_no_ext + f"_{count}_{folder.value}.jpg")

            if folder == SaveType.FACES:
                if type(self.face) == Image:
                    self.face.save(save_full_file_name)
            elif folder == SaveType.FOUND_IMAGE:
                if self.label is not None:
                    dir_name, filename = os.path.split(save_full_file_name)
                    label_path = os.path.join(dir_name, self.prompt if self.prompt is not None else "No Prompt", filename)

                    try:
                        os.makedirs(os.path.dirname(label_path))
                    except FileExistsError:
                        pass

                    self.get_masked_image().convert("RGB").save(label_path)
                else:
                    self.get_masked_image().convert("RGB").save(save_full_file_name)
            else:
                self.get_masked_image().convert("RGB").save(save_full_file_name)

            print(f"saving {folder.value} image to {save_full_file_name}")

        def get_org_image(self):
            return Image.open(self.org_image_path)

        def get_masked_image(self, invert=True, get_hole=False):
            def get_mask_or_alpha(mode="RGB"):
                mask_image_alpha = Image.new(mode, self.get_org_image().size, color=(0, 0, 0, 0))
                draw = ImageDraw.Draw(mask_image_alpha, mode)
                draw.polygon(self.segment.segment, fill=(255, 255, 255))
                return mask_image_alpha

            if get_hole:
                return_mask = ImageChops.invert(get_mask_or_alpha("RGBA"))
            elif invert:
                mask_image_alpha = get_mask_or_alpha("RGBA")
                return_mask = self.get_org_image().convert("RGBA")
                return_mask.alpha_composite(ImageChops.invert(mask_image_alpha))
            else:
                mask_image = get_mask_or_alpha()
                return_mask = mask_image
            return return_mask

    @dataclass
    class PersonList:
        persons: list[PersonJob] = field(default_factory=list)
        found_persons: list[PersonJob] = field(default_factory=list)
        unfound_persons: list[PersonJob] = field(default_factory=list)
        group_by_image_path_list: list[list[PersonJob]] = field(default_factory=list)
        generated_images: list[PersonJob] = field(default_factory=list)

        def get_found_persons(self):
            self.found_persons = []
            self.unfound_persons = []

            for loop_person in self.persons:
                if loop_person.found_face:
                    self.found_persons.append(loop_person)
                else:
                    self.unfound_persons.append(loop_person)

        def save_faces(self, save_type: SaveType):
            pool = ThreadPoolExecutor()
            if save_type == SaveType.UNFOUND_IMAGE:
                list_of_persons = self.unfound_persons

            if save_type == SaveType.DEBUG_IMAGE:
                list_of_persons = self.persons

            if save_type == SaveType.FOUND_IMAGE:
                list_of_persons = self.found_persons

            workers = []
            for count, save_person in enumerate(list_of_persons):
                workers.append(pool.submit(save_person.save_debug_image, save_type, count))

            for worker in workers:
                worker.result()

        def get_prompts(self):
            def get_missing():
                if not check_faces and not prompt_by_body:
                    return
                if check_faces:
                    loop_times = range(2)
                else:
                    loop_times = range(1)

                for times in loop_times:
                    for person_count, person in enumerate(self.persons):
                        if check_faces:
                            checker = person.found_face is False and times != 0
                        else:
                            checker = person.track_label is None
                        person.track_label = person.segment.id

                        if checker:  # this means we found a segmentation map but there was no track_label so we dont know what prompt to give it
                            searching_for_box = True
                            current_frame = person.segment.frame_idx
                            before_count = person_count
                            after_count = person_count
                            run_times = 0

                            while searching_for_box:
                                run_times += 1
                                compare_list = []
                                after_count += 1

                                if current_frame != 0:
                                    before_count -= 1

                                    try:
                                        before = self.persons[before_count]
                                        if before.segment.frame_idx < current_frame:
                                            compare_list.append(before)
                                    except IndexError:
                                        pass

                                try:
                                    after = self.persons[after_count]
                                    if after.segment.frame_idx > current_frame:
                                        compare_list.append(after)
                                except IndexError:
                                    pass

                                if compare_list:
                                    for _, compare_person in enumerate(compare_list):
                                        if searching_for_box is False:
                                            continue

                                        if prompt_by_body:
                                            if compare_person.segment.id is None:
                                                continue

                                        bounding_boxes = zip(person.segment.bounding_box,
                                                             compare_person.segment.bounding_box)
                                        all_difference = 0
                                        for box in bounding_boxes:
                                            all_difference += abs(box[0] - box[1])
                                        if all_difference <= 100:
                                            if times == 0:
                                                person.track_label = compare_person.segment.id
                                                if person.track_label > len(func_prompt_list):
                                                    person.prompt = fallback
                                                else:
                                                    person.prompt = func_prompt_list[int(person.track_label) - 1]
                                                    print(person.prompt)
                                            else:
                                                if not compare_person.found_face:
                                                    continue

                                                person.label = compare_person.label
                                                if person.label > len(func_prompt_list):
                                                    person.prompt = fallback
                                                else:
                                                    if person.label+1 > len(func_prompt_list):
                                                        person.prompt = fallback
                                                    else:
                                                        person.prompt = func_prompt_list[int(person.label)]
                                                        person.found_face = True
                                                        print(person.prompt)

                                            searching_for_box = False

                                        elif person.track_label is not None:
                                            if person.track_label > len(func_prompt_list):
                                                person.prompt = fallback
                                            else:
                                                person.prompt = func_prompt_list[int(person.track_label) - 1]
                                                print(person.prompt)

                                if run_times >= 15:
                                    if not searching_for_box:
                                        continue

                                    print(f"searched {run_times} for match and didnt find")
                                    break

                            if person.prompt is None:
                                print(f"None person on frame {person.segment.frame_idx}")
                        elif prompt_by_body:
                            if person.track_label > len(func_prompt_list):
                                person.prompt = fallback
                            else:
                                person.prompt = func_prompt_list[int(person.track_label) - 1]
                                print(person.prompt)

            if check_faces:
                self.persons = facechecker.get_face_identities(self.persons, face_match_sensitivity)
                self.get_found_persons()

                for person in self.found_persons:
                    try:
                        person.prompt = func_prompt_list[int(person.label)]
                    except IndexError:
                        person.prompt = fallback

                get_missing()
                self.get_found_persons()
                if debug:
                    person_list.save_faces(SaveType.UNFOUND_IMAGE)
                    person_list.save_faces(SaveType.FOUND_IMAGE)
                    print("Finished Saving Debug Images")

            elif not prompt_by_body:
                for person in self.persons:
                    person.prompt = func_prompt_list[0]

                self.group_by_image_path()
                return

            get_missing()
            self.group_by_image_path()

        def group_by_image_path(self):
            person_images = {}
            for person in self.persons:
                found_path = person_images.get(person.org_image_path)
                if found_path is None:
                    person_images[person.org_image_path] = [person]
                else:
                    person_images[person.org_image_path].append(person)

            self.group_by_image_path_list = [value for value in person_images.values()]

        def add_to_generate_list(self, generated_class:PersonJob):
            generated_class.generated_times += 1
            self.generated_images.append(generated_class)

    @dataclass
    class SegmentCoord:
        segment: np.array
        frame_idx: int
        bounding_box: field(default_factory=list)
        path: str
        id = None
        dif = None

    @dataclass
    class SegmentListClass:
        segments: field(default_factory=list)

        def sort_segments(self):
            self.segments = sorted(self.segments, key=lambda x: x.path)

    person_list = PersonList()
    segment_list = SegmentListClass(list())
    host, port = str(server).split(":")
    api = webuiapi.WebUIApi(host=host, port=port)
    unfound_masks = []

    def name_cleaner(name_to_clean:str):
        valid_chars = f"-_.(), {string.ascii_letters}{string.digits}"
        safe_filename = "".join(letter for letter in name_to_clean if letter in valid_chars)
        return safe_filename

    def check_if_masks_exists_and_save_masks(list_of_files: list, current_image_segments=None):
        for _, image_path in enumerate(list_of_files):
            current_image_path = image_path[1] if type(image_path) == tuple else image_path
            not_found = True
            masks_folder = os.path.join(output_path, "masks")
            mask_filename = name_cleaner(current_image_path)
            mask_filename = os.path.join(masks_folder, mask_filename + f".json")

            if not os.path.exists(masks_folder):
                try:
                    os.makedirs(masks_folder)

                except FileExistsError:
                    pass

            else:
                if os.path.exists(mask_filename):
                    with open(mask_filename, "rb") as mask_file:
                        segments_tuple = orjson.loads(mask_file.read())

                    for segment_item in segments_tuple:
                        current_segment = SegmentCoord(segment=np.array(segment_item.get("segment"), dtype=np.float32),
                                                       frame_idx=segment_item.get("frame_idx"),
                                                       bounding_box=segment_item.get("bounding_box"),
                                                       path=segment_item.get("path"))
                        current_segment.id = segment_item.get("id")
                        current_segment.dif = segment_item.get("dif")
                        for value in astuple(current_segment):
                            if value is None and current_segment.frame_idx !=0:
                                break

                        segment_list.segments.append(current_segment)
                        id_set.add(current_segment.id)
                        not_found = False
            if current_image_segments is not None:
                with open(mask_filename, "wb") as mask_file:
                    mask_file.write(orjson.dumps(tuple(current_image_segments), option=orjson.OPT_SERIALIZE_NUMPY))

            if not not_found:
                print(f"found mask for {current_image_path} so no new mask will be generated for it\n")
            else:
                unfound_masks.append(current_image_path)

    check_mask_pool = ThreadPoolExecutor()
    check_mask_workers = []

    for check_image_path in enumerate(glob.glob(source+"\*")):
        check_mask_workers.append(check_mask_pool.submit(check_if_masks_exists_and_save_masks, [check_image_path]))

    for check_mask_worker in check_mask_workers:
        check_mask_worker.result()

    # Load model
    if unfound_masks:
        device = select_device(device)
        is_seg = '-seg' in str(yolo_weights)
        model = DetectMultiBackend(yolo_weights, device=device, dnn=dnn, data=None, fp16=half)
        stride, names, pt = model.stride, model.names, model.pt
        imgsz = check_img_size(imgsz, s=stride)  # check image size
        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
        nr_sources = 1
        segment_list.sort_segments()
        dataset.files = sorted(unfound_masks)
        dataset.nf = len(unfound_masks)
        dataset.count = len(unfound_masks)

        # Create as many strong sort instances as there are video sources
        tracker_list = []
        for i in range(nr_sources):
            tracker = create_tracker(tracking_method, reid_weights, device, half)
            tracker_list.append(tracker, )
            if hasattr(tracker_list[i], 'model'):
                if hasattr(tracker_list[i].model, 'warmup'):
                    tracker_list[i].model.warmup()
        outputs = [None] * nr_sources

        seen, windows, dt = 0, [], (Profile(), Profile(), Profile(), Profile())
        curr_frames, prev_frames = [None] * nr_sources, [None] * nr_sources

    for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):
            with dt[0]:
                im = torch.from_numpy(im).to(device)
                im = im.half() if half else im.float()  # uint8 to fp16/32
                im /= 255.0  # 0 - 255 to 0.0 - 1.0
                if len(im.shape) == 3:
                    im = im[None]  # expand for batch dim

            # Inference
            with dt[1]:

                if is_seg:
                    pred, proto = model(im, augment=augment, visualize=visualize)[:2]
                else:
                    pred = model(im, augment=augment, visualize=visualize)

            # Apply NMS
            with dt[2]:
                if is_seg:
                    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det, nm=32)
                else:
                    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

            # Process detections
            for i, det in enumerate(pred):  # detections per image
                seen += 1
                p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
                p = Path(p)  # to Path
                curr_frames[i] = im0

                s += '%gx%g ' % im.shape[2:]  # print string
                imc = im0.copy() if save_crop else im0  # for save_crop

                if hasattr(tracker_list[i], 'tracker') and hasattr(tracker_list[i].tracker, 'camera_update'):
                    if prev_frames[i] is not None and curr_frames[i] is not None:  # camera motion compensation
                        tracker_list[i].tracker.camera_update(prev_frames[i], curr_frames[i])

                if det is not None and len(det):
                    if is_seg:
                        # scale bbox first the crop masks
                        if retina_masks:
                            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()  # rescale boxes to im0 size
                            masks = process_mask_native(proto[i], det[:, 6:], det[:, :4], im0.shape[:2])  # HWC
                        else:
                            masks = process_mask(proto[i], det[:, 6:], det[:, :4], im.shape[2:], upsample=True)  # HWC
                            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()  # rescale boxes to im0 size

                        # Segments
                        segments = [
                            scale_segments(im0.shape if retina_masks else im.shape[2:], x, im0.shape, normalize=False)
                            for x in reversed(masks2segments(masks))]

                        current_image_segments = []
                        for j, segment in enumerate(segments):
                            segmentation_map = segment

                            # Find the minimum and maximum x and y values
                            min_x = min(x for x, y in segmentation_map)
                            max_x = max(x for x, y in segmentation_map)
                            min_y = min(y for x, y in segmentation_map)
                            max_y = max(y for x, y in segmentation_map)
                            area = (max_x - min_x) * (max_y - min_y)

                            if area <= max_person_size:
                                print(f"Skipping detected person on frame {frame_idx} as they are below the size limit of {org_max_person_size}")
                                continue

                            # Print the bounding box coordinates
                            current_segment = SegmentCoord(segment=segment, frame_idx=frame_idx, bounding_box=(min_x, min_y, max_x, max_y), path=path)
                            segment_list.segments.append(current_segment)
                            current_image_segments.append(current_segment)
                            if frame_idx == 0:
                                current_segment.id = None

                    else:
                        det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()  # rescale boxes to im0 size

                    # Print results
                    for c in det[:, 5].unique():
                        n = (det[:, 5] == c).sum()  # detections per class
                        s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                    # pass detections to strongsort
                    with dt[3]:
                        outputs[i] = tracker_list[i].update(det.cpu(), im0)

                    # draw boxes for visualization

                    if len(outputs[i]) > 0:
                        for j, (output) in enumerate(outputs[i]):
                            bboxes = output[0:4]
                            id = output[4]
                            cls = output[5]
                            conf = output[6]

                            for a, segment in enumerate(current_image_segments):
                                left_dif = abs(segment.bounding_box[0] - bboxes[0])
                                top_dif = abs(segment.bounding_box[1] - bboxes[1])
                                bot_dif = abs(segment.bounding_box[2] - bboxes[2])
                                right_dif = abs(segment.bounding_box[3] - bboxes[3])
                                all_dif = left_dif + top_dif + bot_dif + right_dif
                                segment.dif = all_dif

                            smallest = None
                            for segment in current_image_segments:
                                if smallest is None:
                                    smallest = segment.dif
                                else:
                                    if segment.dif < smallest:
                                        smallest = segment.dif

                            for segment in current_image_segments:
                                if smallest == segment.dif:
                                    segment.id = id
                                    id_set.add(id)

                    if current_image_segments:
                        check_if_masks_exists_and_save_masks([path], current_image_segments)
                else:
                    pass

            try:
                LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{sum([dt.dt for dt in dt]) * 1E3:.1f}ms")
            except AttributeError:
                print(LOGGER.info(f"{s} (no detections)"))

    try:
        del model
        del dt
        del pred
        torch.cuda.empty_cache()
        gc.collect()
    except UnboundLocalError:
        print("no model to delete")
        pass

    id_list = [id for id in id_set if id is not None]
    id_dict = {}

    for count, id_to_change in enumerate(id_list, 1):  #normalizing ids so that we can use them as indexes later
        if id_to_change is None:
            continue

        id_dict.update({id_to_change: count})

    for segment in segment_list.segments:
        person = PersonJob(segment.path)

        person.segment = segment

        person.segment.id = id_dict[person.segment.id] if segment.id is not None else None

        if debug or check_faces:
            person.masked_image = True

        person_list.persons.append(person)

    person_list.persons = sorted(person_list.persons, key=lambda x:x.org_image_path)
    segment_list.sort_segments()
    person_list.get_prompts() #assigns prompts to people and groups images by the path to the file
    ssim_persons = []
    last_generated_image_for_injection = None

    def save_generated_img(image: Image, full_path_to_save: str):
        save_img_dir = os.path.dirname(full_path_to_save)
        if not os.path.exists(save_img_dir):
            os.makedirs(save_img_dir)
        try:
            try:
                image.save(full_path_to_save)
            except PermissionError:
                time.sleep(.1)
                save_generated_img(image, full_path_to_save)
                print(f"Error Saving: {full_path_to_save} retrying")
        except OSError:
            pass

        print(f"Saved image to {full_path_to_save}")
        ssim_persons.append({"image": image, "path": full_path_to_save})
        func_last_generated_image_for_injection = image
        return func_last_generated_image_for_injection

    def get_similar_frames(files_to_check):
        func_skip_set = set()
        for count, file in enumerate(files_to_check):
            if file in func_skip_set:
                continue

            skip_list = []
            frames_count = 0

            current_image = Image.open(file)
            for sub_count, sub_file in enumerate(files_to_check):
                if frames_count >= max_frame_skip:
                    skip_list.pop(-1)
                    for pathfile in skip_list:
                        func_skip_set.add(pathfile)
                    break

                if sub_count <= count:
                    continue

                next_image = Image.open(sub_file)
                sub_similar_ratio = SSIM_PIL.compare_ssim(current_image, next_image)

                if sub_similar_ratio*100 > skip_ratio:
                    frames_count += 1
                    skip_list.append(sub_file)
                else:
                    if skip_list:
                        skip_list.pop(-1)
                        for pathfile in skip_list:
                            func_skip_set.add(pathfile)
                    break
        return func_skip_set

    def save_video(path_to_images, framerate, filename, generated_images_folder=""):

        filename = name_cleaner(filename)
        if not framerate:
            framerate = 30

        non_generated_images = glob.glob(path_to_images+"\*.png")
        generated_images = glob.glob(path_to_images+"\generated_imgs\*.png")

        if generated_images_folder:
            generated_images_folder = glob.glob(generated_images_folder+"\*.png")

        if not generated_images_folder:
            images = non_generated_images
        else:
            images = sorted(non_generated_images+generated_images_folder, key=lambda x: os.path.basename(x))

        sequence = ImageSequenceClip(images, framerate)
        sequence_save_folder = os.path.join(output_path, "videos")

        if not os.path.exists(sequence_save_folder):
            os.makedirs(sequence_save_folder)

        save_filename = os.path.join(sequence_save_folder, filename+f" {random.random()}.mp4")

        sequence.write_videofile(save_filename, framerate, audio=audio_source)

    def skip_list_split(list_to_split, sections=5):
        list_length = len(list_to_split)
        section_size = list_length // sections
        lists = []
        loop_range = range(sections)
        loop_range_size = len(loop_range)
        for section in loop_range:
            section += 1
            slice_amt = section * section_size

            if section == loop_range_size:
                lists.append(list_to_split[slice_amt - section_size:])
            else:
                lists.append(list_to_split[slice_amt - section_size:slice_amt + max_frame_skip + 1])
        return tuple(lists)

    skip_set = set()
    files_in_source_dir = list(glob.glob(source + "\*.png"))

    if frame_gen_auto:
        print(f"Checking for frames to skip using a ratio of {skip_ratio}")
        skip_workers = []
        skip_pool = ThreadPoolExecutor()
        skip_lists = skip_list_split(files_in_source_dir, 6)

        for skip_list in skip_lists:
            res = skip_pool.submit(get_similar_frames, skip_list)
            skip_workers.append(res)

        for skip_worker in skip_workers:
            skip_set.update(skip_worker.result())

        print(f"Found {len(skip_set)} similar image(s) to skip")

    total_jobs = 0
    for count_job in person_list.group_by_image_path_list:
        total_jobs += len(count_job)
    total_jobs = total_jobs-len(skip_set)
    org_job_length = total_jobs-len(skip_set)
    done_jobs = 0
    generated_list = {}
    last_time = 0
    for personjob_list_count, personjob_list in enumerate(person_list.group_by_image_path_list):

        if frame_gen_all or frame_gen_even:
            if personjob_list_count % 2 != 0:
                continue

        ssim_loop = range(ssim_retry) if personjob_list_count != 0 else range(1)
        best_ssim = 0
        variation = 0
        skip_this_frame = False
        #hack need to change later

        if skip_set:
            for set_personjob in personjob_list:
                if set_personjob.org_image_path in skip_set:
                    skip_this_frame = True
                    break

        if skip_if_exists:
            for check_personjob in personjob_list:
                if os.path.exists(check_personjob.full_save_path_with_ext):
                    print(f"file: {check_personjob.full_save_path_with_ext} exists so it will be skipped. And no new image will be generated")
                    skip_this_frame = True
                    break

        if skip_this_frame:
            done_jobs += len(personjob_list)
            remaining_jobs = total_jobs-done_jobs
            print(f"{remaining_jobs} jobs remaining")

            for personjob in personjob_list:
                if generated_list.get(personjob.prompt) is None:
                    generated_list[personjob.prompt] = 1
                else:
                    generated_list[personjob.prompt] += 1
            continue

        for ssim_amount in ssim_loop:
            ssim_amount += 1
            doing_ssim = ssim_amount > 1

            if doing_ssim and variation_scaling != 0:
                variation = variation_scaling

            if doing_ssim and len(ssim_persons) > 1:
                main_image = ssim_persons[0]["image"]
                last_generated_image_to_compare = ssim_persons[-1]["image"]
                ssim_difference = SSIM_PIL.compare_ssim(main_image, last_generated_image_to_compare)
                print(f"Compared Image 1 to Image {len(ssim_persons)} and they were {round(ssim_difference, 3)*100}% similar")

                if ssim_difference >= ssim_threshold/100:
                    print(f"This is above the similarity threhsold of {ssim_threshold}% that was set so this Image has been selected")
                    last_generated_image_for_injection = save_generated_img(last_generated_image_to_compare, ssim_persons[-1]["path"])

                    done_jobs += 1
                    remaining_jobs = total_jobs - done_jobs
                    print(f"{remaining_jobs} jobs remaining")

                    ssim_persons = []
                    break
                else:
                    print(f"This is below the similarity threhsold of {ssim_threshold} that was set so will keep searching {ssim_amount}/{ssim_retry} times")
                    done_jobs -= 1
                    if ssim_difference > best_ssim:
                        best_ssim = ssim_difference
                        best_ssim_image = last_generated_image_to_compare
                        best_ssim_image_path = ssim_persons[-1]["path"]

            if ssim_amount == ssim_retry and doing_ssim:
                print(f"Reached search limit {ssim_amount}/{ssim_retry} limit choosing most similar image in batch which was {round(best_ssim, 3)*100}% similar")
                last_generated_image_for_injection = save_generated_img(best_ssim_image, best_ssim_image_path)

                done_jobs += 1
                remaining_jobs = total_jobs - done_jobs
                print(f"{remaining_jobs} jobs remaining")

                ssim_persons = []
                break

            completed_images = []

            for person_count, personjob in enumerate(personjob_list, 1):
                if personjob.prompt is None:
                    if ignore_if_no_prompt:
                        continue
                    else:
                        completed_images = []
                        for sub_personjob in personjob_list:
                            sub_personjob.generated_times = 0
                        break

                for loop_times in range(loopback):
                    loop_times = loop_times
                    loop_times += 1
                    print_prompt = personjob.prompt if len(personjob.prompt) < 60 else personjob.prompt[:60]+"..."

                    if loop_times > 1:
                        print(f"generating prompt: '{print_prompt}' for person {person_count} again since loopback is {loopback} current: {loop_times}/{loopback}")

                    if loopback > 1 and loop_times > 1 and person_list.generated_images:
                        if not use_control_net:
                            gen_image = image.images[-1]
                        else:
                            gen_image = image.images[0]

                    elif inject_last_image and last_generated_image_for_injection is not None:
                        last_generated_image = last_generated_image_for_injection
                        last_generated_image = last_generated_image.filter(ImageFilter.GaussianBlur(inject_last_image["Blur"]))
                        mask = ImageChops.invert(personjob.get_masked_image(get_hole=True))
                        last_generated_image = Image.composite(last_generated_image, mask, mask).convert("RGBA")
                        last_generated_image_array = np.array(last_generated_image)
                        opaque = last_generated_image_array[:, :, 3] == 255
                        opacity = inject_last_image["Opacity"]
                        last_generated_image_array[opaque, 3] = opacity * 255
                        last_generated_image = Image.fromarray(last_generated_image_array).convert("RGBA")
                        blended_image = Image.alpha_composite(personjob.get_org_image().convert("RGBA"), last_generated_image)
                        gen_image = blended_image
                        print(f"Injecting previous image into generation at blur {inject_last_image['Blur']} and opacity {inject_last_image['Opacity']}")

                    else:
                        gen_image = personjob.get_org_image()

                    if loop_times <= 1:
                        print(f"generating prompt: '{print_prompt}' for person {person_count}")
                    api_ssim_amount = ssim_amount if seed_scaling else 0

                    if pipe.poll():
                        instruction = pipe.recv()
                        if instruction == "STOP":
                            return False

                    prompt = personjob.prompt
                    gen_times = generated_list.get(prompt)

                    transition_frames = 30
                    mid_transition_frames = 15

                    if loop_times <= 1:
                        print(f"generating prompt: '{print_prompt}' for person {person_count}")

                    if last_time != 0:
                        cur_time = time.perf_counter()
                        runtime = cur_time-last_time
                        print(runtime)
                        total_time = round(runtime*total_jobs)
                        print("")
                        print(f"estimated time remaining {total_time} seconds")
                        print("")

                    last_time = time.perf_counter()
                    image = api.img2img(images=[gen_image],
                                        prompt=prompt,
                                        mask_image=personjob.get_masked_image(False),
                                        inpainting_fill=1,
                                        cfg_scale=cfg,
                                        denoising_strength=denoising,
                                        inpaint_full_res=True,
                                        steps=steps,
                                        sampler_index=sampler,
                                        restore_faces=restore_faces,
                                        negative_prompt=negative_prompt,
                                        height=height,
                                        width=width,
                                        seed=seed+loop_times+api_ssim_amount,
                                        subseed=ssim_amount,
                                        subseed_strength=variation,
                                        use_control_net=use_control_net,
                                        controlnet_model=control_net_model,
                                        controlnet_module=control_net_processor,
                                        controlnet_lowvram=control_net_lowvram
                                        )

                    person_list.add_to_generate_list(personjob)

                    if loop_times == loopback:

                        done_jobs += 1
                        remaining_jobs = total_jobs - done_jobs
                        print(f"{remaining_jobs} jobs remaining")

                        if use_control_net:
                            image.images = [image.images[0]]

                        completed_images.append([image.images[-1], personjob.get_masked_image(get_hole=True)])

            if len(completed_images) == 1:
                last_generated_image_for_injection = save_generated_img(completed_images[0][0], personjob_list[0].full_save_path_with_ext)

            elif len(completed_images) > 1:
                final_image = functools.reduce(lambda x, y: Image.composite(x[0] if type(x) is list else x, y[0], y[1]), completed_images)
                last_generated_image_for_injection = save_generated_img(final_image, personjob_list[0].full_save_path_with_ext)
            print("")

    generated_images_save_folder = os.path.join(save_folder, "generated_imgs")
    if frame_gen_auto or frame_gen_even:

        if frame_gen_auto:
            print(f"Generating {len(skip_set)} skipped frames")

        RIFE.frame_generator(path_of_images=save_folder, save_folder=generated_images_save_folder)
        save_video(path_to_images=save_folder, filename=name_prompt, framerate=fps, generated_images_folder=generated_images_save_folder)

    elif frame_gen_all:
        RIFE.frame_generator(path_of_images=save_folder, save_folder=generated_images_save_folder)
        RIFE.frame_generator(path_of_images=generated_images_save_folder, save_folder=generated_images_save_folder)
        save_video(path_to_images=generated_images_save_folder, filename=name_prompt, framerate=fps)

    else:
        save_video(path_to_images=save_folder, filename=name_prompt, framerate=fps)

    print(f"Saving video to {os.path.join(output_path, 'videos')}")
    return True


import shutil
import time
import urllib.parse
from dataclasses import dataclass
from typing import Callable, Dict

import dateutil.parser
import requests
import rich.progress as progress
import urllib3
from plib import Path
from retry import retry
from tqdm.utils import CallbackIOWrapper

from .progress import UIProgress

TRIES = 5

@dataclass
class Downloader:
    url: str
    dest: Path | str | None = None
    folder: Path | str = None
    session: requests.Session | None = None
    headers: Dict | None = None
    retry_amount: int = -1
    timeout: int = 10
    progress_callback: Callable = lambda p: None
    skip_same_size: bool = False
    set_time: bool = True  # set modified time to modified time on server

    def __post_init__(self):
        self.dest = self.prepare_dest(self.dest, self.folder)
        self.session = self.prepare_session(self.session, self.headers)

    def prepare_session(self, session, headers):
        session = session or requests.Session()
        browser = (
            "Mozilla/5.0 (X11; Linux x86_64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) "
            "Chrome/85.0.4183.83 Safari/537.36"
        )
        download_time = time.gmtime(self.dest.mtime)
        download_time = time.strftime("%a, %d %b %Y %H:%M:%S GMT", download_time)
        headers = (headers or {}) | {
            "User-Agent": browser,
            "If-Modified-Since": download_time,
        }
        session.headers.update(headers)
        return session

    def prepare_dest(self, dest, folder):
        if dest is None:
            url_path = urllib.parse.urlparse(self.url).path
            name = urllib.parse.unquote(url_path).split("/")[-1]
            dest = name or "download"

        dest = Path(folder) / dest if folder else Path(dest)
        return dest

    @retry(
        (requests.exceptions.RequestException, urllib3.exceptions.HTTPError),
        tries=TRIES,
    )
    def download(self):
        self.retry_amount += 1
        headers = {"Range": f"bytes={self.temp_dest.size}-"}
        if "If-Modified-Since" in self.session.headers:
            self.session.headers.pop("If-Modified-Since")
        stream = self.session.get(
            self.url, timeout=self.timeout, stream=True, headers=headers
        )
        if stream.status_code != 304:
            self.download_modified(stream)

    def download_modified(self, stream):
        if stream.status_code == 416:  # range not satifiable or supported
            stream = self.session.get(self.url, timeout=self.timeout, stream=True)
        elif "Content-Range" in stream.headers:
            start = int(
                stream.headers["Content-Range"].split("bytes ")[1].split("-")[0]
            )
            if self.temp_dest.size > start:
                self.temp_dest.unlink()
                stream = self.session.get(self.url, timeout=self.timeout, stream=True)
        else:
            assert stream.ok

        skip = self.skip_same_size and self.dest.size == int(
            stream.headers["Content-Length"]
        )
        if not skip:
            self.start_download(stream)

    def start_download(self, stream):
        total = (
            int(
                stream.headers["Content-Length"]
            )  # length that still needs to be received
            if "Content-Length" in stream.headers
            else int(stream.headers["Content-Range"].split("/")[-1])
        )

        progress = UIProgress(self.description, total=total + self.dest.size)
        # downloaded size added to todo and done

        def progres_callback(value):
            progress.advance(value)
            self.progress_callback(value / total)

        progres_callback(self.dest.size)  # already downloaded part is progress

        stream_raw = CallbackIOWrapper(progres_callback, stream.raw)

        chunk_size = self.truncate(total // 10, 1 * 1024, 128 * 1024)
        with progress, self.temp_dest.open("ab") as fp:
            shutil.copyfileobj(stream_raw, fp, length=chunk_size)

        self.temp_dest.rename(self.dest)
        modified_key = "last-modified"
        if self.set_time and modified_key in stream.headers:
            server_mtime = stream.headers[modified_key]
            self.dest.mtime = dateutil.parser.parse(server_mtime).timestamp()

    @property
    def temp_dest(self):
        return self.dest.with_suffix(self.dest.suffix + ".part")

    @property
    def description(self):
        return (
            f"[retry {self.retry_amount}/{TRIES - 1}] {self.dest.name}"
            if self.retry_amount > 0
            else self.dest.name
        )

    @staticmethod
    def truncate(value, min_value, max_value):
        if value < min_value:
            value = min_value
        elif value > max_value:
            value = max_value
        return value