import gzip
from timeit import timeit
import numpy as np
import pandas as pd

from tf.core.data import Data
from tf.core.timestamp import Timestamp
from tf.advanced.helpers import dm

from pack import deepSize

TEMP = "_temp"
FEATURES = """
    vs
    g_word_utf8
    rela
""".strip().split()

HEAD = f"""\
feature | length | start | end | NaN | %NaN | python dict | numpy | pandas | pandas-sparse
---     | ---    | ---   | --- | --- | ---  | --- | ---  | --- | ---
"""

HEAD_TIME = f"""\
feature | python dict | numpy | pandas | pandas-sparse
---     | ---  | ---   | --- | ---
"""

NONES = {None, "NA"}

TIMES = 1
KEYS = (100001, 1000001)

T = Timestamp()

class FeatureTest:
    def __init__(self, feat):
        dataObj = Data(f"{TEMP}/{feat}.tf", T)
        dataObj.load()
        T.indent(level=0)
        data  = dataObj.data
        start = min(data)
        end = max(data)
        self.start = start
        self.end = end
        self.nan = sum(1 for n in range(start, end + 1) if data.get(n, None) in NONES)
        self.feat = feat
        self.data  = data
        self.ln = len(data)
        self.mem = deepSize(data) // (1024 * 1024)
        
    def adjust(self, totalMax):
        self.totalMax = totalMax
        self.nanPerc = self.nan * 100 / totalMax
        
    def numpy(self):
        data = self.data
        totalMax = self.totalMax
        array = [data.get(i, "") for i in range(totalMax + 1)]
        dataN = np.array(array, np.str)
        self.dataN = dataN
        self.memN = deepSize(dataN) // (1024 * 1024)
        
    def pandas(self):
        data = self.data
        dataP = pd.Series(data, dtype="string")
        self.dataP = dataP
        self.memP = dataP.memory_usage(index=True, deep=True) // (1024 * 1024)

        totalMax = self.totalMax
        array = [data.get(i, "NA") or "NA" for i in range(totalMax + 1)]
        dataSP = pd.Series(pd.arrays.SparseArray(array, fill_value="NA", dtype="string"))
        self.dataSP = dataSP
        self.memSP = dataSP.memory_usage(index=True, deep=False) // (1024 * 1024)
        
    def accessTime(self, times):
        data = self.data
        locs = locals()
        self.access = sum(timeit(f"data.get({key}, None)", globals=locs, number=times) for key in KEYS)
        
    def accessTimeN(self, times):
        data = self.dataN
        locs = locals()
        self.accessN = sum(
            timeit(f"data[{key}] if {key} < data.size else None", globals=locs, number=times)
            for key in KEYS
        )
        
    def accessTimeP(self, times):
        dataP = self.dataP
        locs = locals()
        self.accessP = sum(timeit(f"dataP.get({key})", globals=locs, number=times) for key in KEYS)
        
    def accessTimeSP(self, times):
        dataSP = self.dataSP
        locs = locals()
        self.accessSP = sum(timeit(f"dataSP.get({key})", globals=locs, number=times) for key in KEYS)
        
    def report(self):
        return (
            f"{self.feat} | "
            f"{self.ln} | "
            f"{self.start} | "
            f"{self.end} | "
            f"{self.nan} | "
            f"{self.nanPerc} | "
            f"{self.mem} MB | "
            f"{self.memN} MB | "
            f"{self.memP} MB | "
            f"{self.memSP} MB\n"
        )
    
    def reportTime(self):
        return (
            f"{self.feat} | "
            f"{self.access} s | "
            f"{self.accessN} s | "
            f"{self.accessP} s | "
            f"{self.accessSP} s\n"
        )

class DataTest:
    def __init__(self):
        totalMax = 0
        features = {}
        T.indent(reset=True)
        
        for feat in FEATURES:
            T.info(f"stage1 {feat}")
            fObj = FeatureTest(feat)
            if fObj.end > totalMax:
                totalMax = fObj.end
            features[feat] = fObj
        T.info("done")
            
        self.features = features
        
        for (feat, fObj) in features.items():
            T.info(f"stage2 {feat}")
            fObj.adjust(totalMax)
            fObj.numpy()
            fObj.pandas()
        T.info("done")
            
    def accessTime(self, times):
        features = self.features
        for (feat, fObj) in features.items():
            T.info(f"timing {feat}")
            fObj.accessTime(times)
            fObj.accessTimeN(times)
            fObj.accessTimeP(times)
            fObj.accessTimeSP(times)
        
    def report(self):
        features = self.features
        
        md = HEAD
        
        for fObj in features.values():
            md += fObj.report()
            
        dm(md)

    def reportTime(self):
        features = self.features
        
        md = HEAD_TIME
        
        for fObj in features.values():
            md += fObj.reportTime()
            
        dm(md)
        
    def test(self, feat, nodes):
        dt = self.features[feat]

        md = f"""\
tf | numpy | pandas | pandas-sparse
---|---|---|---
"""

        for i in nodes:
            md += f"{dt.data[i]} | {dt.dataN[i]} | {dt.dataP[i]} | {dt.dataSP[i]}\n"

        dm(md)