import json
import resource
import platform
import sys
from pathlib import Path
from typing import *

import submitit
from tqdm import tqdm
import time

from collections import Counter

# Posix based file locking (Linux, Ubuntu, MacOS, etc.)
#   Only allows locking on writable files, might cause
#   strange results for reading.
import fcntl, os


def track_jobs_with_pbar(jobs):
    num_completed = 0
    with tqdm(total=len(jobs)) as pbar:
        while any(job.state not in ["COMPLETED", "FAILED", "DONE"] for job in jobs):
            time.sleep(2)
            state2count = Counter([j.state for j in jobs])
            newly_completed = state2count["COMPLETED"] - num_completed
            pbar.update(newly_completed)
            num_completed = state2count["COMPLETED"]
            s = [f"{k}: {v}" for k, v in state2count.items()]
            pbar.set_description(" | ".join(s))


def memory_limit(percentage: float):
    """
    只在linux操作系统起作用
    """
    if platform.system() != "Linux":
        print('Only works on linux!')
        return
    soft, hard = resource.getrlimit(resource.RLIMIT_AS)
    resource.setrlimit(resource.RLIMIT_AS, (int(get_memory() * 1024 * percentage), hard))


def get_memory():
    with open('/proc/meminfo', 'r') as mem:
        free_memory = 0
        for i in mem:
            sline = i.split()
            if str(sline[0]) in ('MemFree:', 'Buffers:', 'Cached:'):
                free_memory += int(sline[1])
    return free_memory


def memory(percentage=0.8):
    def decorator(function):
        def wrapper(*args, **kwargs):
            memory_limit(percentage)
            try:
                function(*args, **kwargs)
            except MemoryError:
                mem = get_memory() / 1024 / 1024
                print('Remain: %.2f GB' % mem)
                sys.stderr.write('\n\nERROR: Memory Exception\n')
                sys.exit(1)

        return wrapper

    return decorator


def log_hparam_results(writer, model_settings, metric_dict):
    hparam_dict = {k: v for k, v in model_settings.items() if k != 'device' and k != 'in_features'}
    writer.add_hparams(hparam_dict=hparam_dict,
                       metric_dict=metric_dict)


def lock_file(f):
    if f.writable(): fcntl.lockf(f, fcntl.LOCK_EX)


def unlock_file(f):
    if f.writable(): fcntl.lockf(f, fcntl.LOCK_UN)


class AtomicOpen:
    """
    Class for ensuring that all file operations are atomic, treat
    initialization like a standard call to 'open' that happens to be atomic.
    This file opener *must* be used in a "with" block.
    https://stackoverflow.com/questions/489861/locking-a-file-in-python
    """

    # Open the file with arguments provided by user. Then acquire
    # a lock on that file object (WARNING: Advisory locking).
    def __init__(self, path, *args, **kwargs):
        # Open the file and acquire a lock on the file before operating
        self.file = open(path, *args, **kwargs)
        # Lock the opened file
        lock_file(self.file)

    # Return the opened file object (knowing a lock has been obtained).
    def __enter__(self, *args, **kwargs): return self.file

    # Unlock the file and close the file object.
    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
        # Flush to make sure all buffered contents are written to file.
        self.file.flush()
        os.fsync(self.file.fileno())
        # Release th


def save_json(loss_dict_combined, model_settings, prefix=""):
    json_dict = json.dumps(loss_dict_combined, indent=4)
    json_dir = Path(f"{prefix}{model_settings['experiment']}_jsons")  # .resolve()
    os.makedirs(json_dir, exist_ok=True)
    json_path = json_dir / f"{str(model_settings['sampling_method'])}_{model_settings['n_samples_constraints']:05}_{model_settings['pde_weight']:.0f}_{model_settings['seed']:04}_run{model_settings['run_num']:05}.json"
    with open(json_path, "w") as outfile:
        outfile.write(json_dict)


def clean_up_jobs(jobs: List[submitit.Job], pickle_name="run"):
    outputs = []
    for job in jobs:
        try:
            outputs.append(job.result())
        except Exception as e:
            try:
                outputs.append(str(e))
            except Exception as note:
                print(e)
    import pickle

    try:
        with open(f"{pickle_name}.pickle", 'wb') as filehandler:
            pickle.dump(outputs, filehandler)
    except Exception as e:
        print("failed dumping pickle")

    try:
        with open(f"{pickle_name}_jobs.pickle", 'wb') as filehandler:
            pickle.dump(jobs, filehandler)
        return True
    except Exception as e:
        print("failed dumping pickle")
        return False
