# #  -*- coding: utf-8 -*-
# #
# #  Copyright (c) 2024 Featrix, Inc, All Rights Reserved
# #
# #  Proprietary and Confidential.  Unauthorized use, copying or dissemination
# #  of these materials is strictly prohibited.
# #
# import time

# ============================================================================
# STUB CLASSES FOR BACKWARD COMPATIBILITY
# ============================================================================
# These stub classes exist solely to allow old training checkpoints to be
# loaded successfully. When PyTorch saves a checkpoint using torch.save(),
# it pickles the entire model state, including any StopWatch, Interval, or
# TimedIterator objects that were embedded in the model or training state.
#
# When loading old checkpoints, pickle needs these classes to exist in the
# module namespace to successfully unpickle the saved objects. Without these
# stubs, loading old checkpoints fails with:
#   "Can't get attribute 'StopWatch' on <module 'featrix.neural.stopwatch'>"
#
# These classes are intentionally minimal stubs - they don't implement the
# original functionality since StopWatch is no longer actively used in the
# codebase. They only need to be unpicklable, not functional.
# ============================================================================

class StopWatch:
    """Stub class for backward compatibility with old checkpoints."""
    def __init__(self, *args, **kwargs):
        # Accept any arguments to allow unpickling of old objects
        # Initialize default attributes if not already set (e.g., during unpickling)
        if not hasattr(self, 'event_timestamps'):
            self.event_timestamps = {}
        if not hasattr(self, 'intervals'):
            self.intervals = set()
        if not hasattr(self, 'state'):
            self.state = 'ready'
    
    def __getstate__(self):
        # Allow pickling
        return self.__dict__
    
    def __setstate__(self, state):
        # Allow unpickling - restore all attributes from the pickled state
        self.__dict__.update(state)
        # Ensure required attributes exist even if not in state
        if 'event_timestamps' not in self.__dict__:
            self.event_timestamps = {}
        if 'intervals' not in self.__dict__:
            self.intervals = set()
        if 'state' not in self.__dict__:
            self.state = 'ready'


class Interval:
    """Stub class for backward compatibility with old checkpoints."""
    def __init__(self, *args, **kwargs):
        # Accept any arguments to allow unpickling of old objects
        pass
    
    def __getstate__(self):
        return self.__dict__
    
    def __setstate__(self, state):
        self.__dict__.update(state)


class TimedIterator:
    """Stub class for backward compatibility with old checkpoints."""
    def __init__(self, *args, **kwargs):
        # Accept any arguments to allow unpickling of old objects
        pass
    
    def __getstate__(self):
        return self.__dict__
    
    def __setstate__(self, state):
        self.__dict__.update(state)


# class StopWatch:
#     """
#     `start` and `end` are provided to help with synchronization. We want to make sure that all ticks
#     come from the same cycle. If we had a single dictionary of timestamps that was updated continuously,
#     it would be possible to end up in an inconsistent state where some timestamps came from the current
#     cycle while others came from the previous cycle.
#     """

#     start_event = "<timer start>"
#     stop_event = "<timer stop>"
#     # The interval that spans between "start" and "stop" events
#     total_interval_name = "total"

#     def __init__(self):
#         self.reset()

#     def start(self):
#         # Make sure the stopwatch was properly stopped.
#         # Failing here helps catch errors where the timer is started
#         # twice, indicating possible confusion about where the cycle starts.
#         assert self.state == "ready"

#         self.reset()

#         self.state = "running"

#         self.tick(self.start_event)

#     def stop(self):
#         self.tick(self.stop_event)
#         self.state = "ready"

#     def reset(self):
#         self.event_timestamps = dict()
#         self.intervals = set()
#         self.state = "ready"

#     def tick(self, event):
#         # Ticks can only be registered if the stopwatch is running.
#         assert self.state == "running"
#         if event in self.event_timestamps:
#             raise ValueError(f"Cannot register event {event} more than once.")

#         self.event_timestamps[event] = time.perf_counter()

#     @classmethod
#     def new_from_state(cls, state_dict):
#         obj = cls()
#         obj.event_timestamps = state_dict

#         return obj

#     def interval(self, name):
#         interval = Interval(self, name)
#         self.intervals.add(name)
#         return interval

#     def get_interval_durations(self):
#         durations = dict()

#         # get the duration of the interval between start and stop
#         durations[self.total_interval_name] = (
#             self.event_timestamps[self.stop_event]
#             - self.event_timestamps[self.start_event]
#         )

#         for interval_name in self.intervals:
#             duration = self.interval(interval_name).duration()
#             durations[interval_name] = duration

#         return durations


# class Interval:
#     """A convenient context-manager-like interface to register a timing interval with a StopWatch object.

#     The interval is recorded as two events - start and stop.
#     """

#     start_suffix = "-<featrix-start>"
#     end_suffix = "-<featrix-end>"

#     def __init__(self, stopwatch: StopWatch, name: str):
#         self.stopwatch = stopwatch
#         self.name = name

#     def __enter__(self):
#         self.stopwatch.tick(self.name + self.start_suffix)

#     def __exit__(self, *args):
#         self.stopwatch.tick(self.name + self.end_suffix)

#     def duration(self):
#         start_timestamp = self.stopwatch.event_timestamps[self.name + self.start_suffix]
#         end_timestamp = self.stopwatch.event_timestamps[self.name + self.end_suffix]

#         return end_timestamp - start_timestamp


# class TimedIterator:
#     def __init__(self, iterable):
#         self.iterable = iterable
#         self.stopwatch = StopWatch()

#     def __iter__(self):
#         self.iterator = iter(self.iterable)
#         return self

#     def __next__(self):
#         self.stopwatch.start()

#         # No named interval needed - the timing will be captured in the default interval
#         # between start and stop.
#         try:
#             result = next(self.iterator)
#         finally:
#             # StopIterator is thrown when there's no more elements to iterate over,
#             # and we need to stop the stopwatch even if exception is thrown.
#             self.stopwatch.stop()

#         return result
