import psutil
import time
import gc
from cpumonitor import CPUMonitor

class Measure:
    def __init__(self, pid, useCpuMonitor=False, cpuMonitorInterval=1.0):
        self.pid = pid
        self.useCpuMonitor = useCpuMonitor
        self.cpuMonitorInterval = cpuMonitorInterval

        self.start_memory = self.get_current_mem_wset()
        self.start_time = time.time()
        # print(f"..Object creation time mem usage: {int(self.start_memory):,} KBs")

    def start(self):
        self.start_time = time.time()
        self.start_memory = self.get_current_mem_wset()
        # print(f"..Starting out mem usage: {int(self.start_memory):,} KBs")
        if self.useCpuMonitor:
            self.monitor = CPUMonitor(self.cpuMonitorInterval)
            self.monitor.start()
        else:
            print("..(not using cpumonitor)")

    def end(self):
        end_time = time.time()
        if self.useCpuMonitor:
            self.monitor.stop()
            print(" ----------------------------------------- ")
            print(f" > Peak CPU usage: {self.monitor.peak_cpu}%'")

        end_memory = self.get_current_mem_wset()
        end_peak_memory = self.get_current_mem_peak_wset()

        timely_mem_usage = end_peak_memory - self.start_memory
        time_taken_sec = (end_time - self.start_time)

        print(f" > Memory surge/usage: {int(timely_mem_usage):,} KBs ({int(self.start_memory):,} KBs -> {int(end_peak_memory):,} KBs)")
        print(f" > Time taken: {time_taken_sec:,} sec")
        print(" ----------------------------------------- ")
        
        # Force garbage collection
        gc.collect()

    def get_current_mem_wset(self):
        return psutil.Process(self.pid).memory_full_info().wset / 1024

    def get_current_mem_peak_wset(self):
        return psutil.Process(self.pid).memory_full_info().peak_wset / 1024
