# Import libraries
import gzip
import json
import os
import numpy as np
import pandas as pd
import datetime

from ipynb.fs.defs.metric_template import decode_run_length
from ipynb.fs.defs.metric_template import CephCommunication
from ipynb.fs.defs.metric_template import save_to_disk, read_from_disk
from dotenv import load_dotenv, find_dotenv

load_dotenv(find_dotenv())

## Specify variables

METRIC_NAME = "correlation"

# Specify the path for input grid data,
INPUT_DATA_PATH = "../../../../data/raw/testgrid_183.json.gz"

# Specify the path for output metric data
OUTPUT_DATA_PATH = f"../../../../data/processed/metrics/{METRIC_NAME}"

## CEPH Bucket variables
## Create a .env file on your local with the correct configs,
s3_endpoint_url = os.getenv("S3_ENDPOINT")
s3_access_key = os.getenv("S3_ACCESS_KEY")
s3_secret_key = os.getenv("S3_SECRET_KEY")
s3_bucket = os.getenv("S3_BUCKET")
s3_path = os.getenv("S3_PROJECT_KEY", "ai4ci/testgrid/metrics")
s3_input_data_path = "raw_data"
AUTOMATION = os.getenv("IN_AUTOMATION")

## Import data
timestamp = datetime.datetime.today()

if AUTOMATION:
    filename = f"testgrid_{timestamp.day}{timestamp.month}.json"
    cc = CephCommunication(s3_endpoint_url, s3_access_key, s3_secret_key, s3_bucket)
    s3_object = cc.s3_resource.Object(s3_bucket, f"{s3_input_data_path}/{filename}")
    file_content = s3_object.get()["Body"].read().decode("utf-8")
    testgrid_data = json.loads(file_content)

else:
    with gzip.open(INPUT_DATA_PATH, "rb") as read_file:
        testgrid_data = json.load(read_file)

failure_groups = []

for tab in list(testgrid_data.keys()):
    for grid in testgrid_data[tab].keys():
        current_grid = testgrid_data[tab][grid]

        tests = [
            current_grid["grid"][i]["name"] for i in range(len(current_grid["grid"]))
        ]
        # unroll the run-length encoding and set bool for flake or not (x==13)
        decoded = [
            (
                np.array(decode_run_length(current_grid["grid"][i]["statuses"])) == 12
            ).tolist()
            for i in range(len(current_grid["grid"]))
        ]

        matrix = pd.DataFrame(zip(tests, decoded), columns=["test", "values"])
        matrix = pd.DataFrame(matrix["values"].to_list(), index=matrix["test"])

        for c, items in matrix.iteritems():
            if len(items[items].index) > 1:
                failure_groups.append(items[items].index)

failure_groups = pd.Series(failure_groups)

len(failure_groups)

vocab = set()
count = 0
for fg in failure_groups:
    count += len(fg)
    vocab.update(fg)

vocab = list(vocab)
print(count)
len(vocab)

len(pd.Series(vocab).unique()) == len(vocab)

def encode_tests(job):
    encoded = []
    for v in vocab:
        if v in job:
            encoded.extend([1])
        else:
            encoded.extend([0])
    return encoded

encoded = failure_groups.apply(encode_tests)

encoded.head()

df_encoded = pd.DataFrame(encoded.array, columns=vocab)
df_encoded.head()

# percent that each test is present in the data; percent failure
perc_present = df_encoded.sum() / len(df_encoded)
perc_present.sort_values(ascending=False).head(3)

# Total failure count present in the data; failure per test
occurrence_count = df_encoded.sum()
occurrence_count.sort_values(ascending=False).head(3)

filter_unique = list(occurrence_count[occurrence_count.values <= 1].index)

df_encoded = df_encoded.drop(filter_unique, axis=1)

df_encoded.shape

# this takes time with full dataset - ~ 2 hours may need to use different approach
# todo try with dask
corr_matrix = df_encoded.corr()

# For each feature, find the other features that are correlated by more than 0.9
top_correlation = {}

for c in corr_matrix.columns:
    top_correlation[c] = []
    series = corr_matrix.loc[c]

    for i, s in enumerate(series):
        if s > 0.90 and series.index[i] != c:
            top_correlation[c].append((series.index[i], s))

len(top_correlation)

# top_correlation has a number of empty rows as not all tests have high correlations with others,
# lets grab only the sets that have at least 1 highly correlated test

pd.set_option("display.max_colwidth", 150)
corr_sets = []
for i in top_correlation.items():
    if len(i[1]) >= 1:
        corr_sets.append(i)
print(f"{len(corr_sets)} sets of correlated tests \n")
print(f"Feature of interest: {corr_sets[1][0]}")
pd.DataFrame(corr_sets[1][1], columns=["test_name", "correlation coefficient"])

if not AUTOMATION:
    test_name = "openshift-tests.[k8s.io] Security Context When creating a container with runAsUser should run the container with uid 65534 [LinuxOnly] [NodeConformance] [Conformance] [Suite:openshift/conformance/parallel/minimal] [Suite:k8s]"  # noqa
    num = occurrence_count.loc[test_name]
    print(f"{num} : the number of times this test failed in our data set")

lst = []
focus = corr_sets[1][1]
for j in focus:
    lst.append((j[0], occurrence_count.loc[j[0]]))

pd.DataFrame(lst, columns=["test_name", "num_occurrences"])

save = pd.DataFrame(corr_sets, columns=["test_name", "correlated_tests"])
save["correlated_tests"] = save["correlated_tests"].apply(str)

if AUTOMATION:
    cc = CephCommunication(s3_endpoint_url, s3_access_key, s3_secret_key, s3_bucket)
    cc.upload_to_ceph(
        save,
        s3_path,
        f"{METRIC_NAME}/{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    )
else:
    save_to_disk(
        save,
        OUTPUT_DATA_PATH,
        f"{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    )

## Sanity check to see if the dataset is the same
if AUTOMATION:
    sanity_check = cc.read_from_ceph(
        s3_path,
        f"{METRIC_NAME}/{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    ).head()
else:
    sanity_check = read_from_disk(
        OUTPUT_DATA_PATH,
        f"{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    ).head()

sanity_check

average_corr = save["correlated_tests"].apply(len).mean()
metric_to_save = pd.DataFrame(
    [[timestamp, average_corr]],
    columns=["timestamp", "average_number_of_correlated_failures"],
)


if AUTOMATION:
    cc.upload_to_ceph(
        metric_to_save,
        s3_path,
        f"avg_{METRIC_NAME}/avg_{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    )
else:
    save_to_disk(
        metric_to_save,
        OUTPUT_DATA_PATH,
        f"avg_{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    )

## Sanity check to see if the dataset is the same

if AUTOMATION:
    sanity_check = cc.read_from_ceph(
        s3_path,
        f"avg_{METRIC_NAME}/avg_{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    ).head()
else:
    sanity_check = read_from_disk(
        OUTPUT_DATA_PATH,
        f"avg_{METRIC_NAME}-{timestamp.year}-{timestamp.month}-{timestamp.day}.parquet",
    ).head()

sanity_check