import pandas as pd
import timesfm
import logging


class Predictions:
    def __init__(self, 
                 context_length: int = 512, 
                 horizon_length: int = 14, 
                 backend: str = "cpu", 
                 checkpoint: str = "google/timesfm-1.0-200m",
                 ) -> None:
        
        logging.info("Initializing Predictions class")

        self.default_step_size = horizon_length

        """
            Initialize timesfm model with the following parameters:
            context_len: The max length of context on which the predictions will be done | Currently supports a max of 512
            horizon_len: The number of future days for which the predictions will be made | reccomended horizon length <= context length
            input_patch_len: Fixed value to load 200m model
            output_patch_len: Fixed value to load 200m model
            num_layers: Fixed value to load 200m model
            model_dims: Fixed value to load 200m model
            backend: The backend to be used for the model | Currently supports "cpu", "gpu", "tpu"
        """
        self.tfm = timesfm.TimesFm(
            context_len = min(context_length, 512), 
            horizon_len=horizon_length,             
            input_patch_len=32,                     
            output_patch_len=128,                   
            num_layers=20,                          
            model_dims=1280,                        
            backend=backend,                        
        )
        
        logging.info("Loading model from checkpoint")
        self.tfm.load_from_checkpoint(repo_id=checkpoint)
        
        logging.info("Model loaded successfully")

    
    def data_preprocess(self, data: pd.DataFrame, target_colm: str, date_colm: str) -> None:
        
        self.data = data.copy()
        self.target_colm = target_colm
        self.default_window_size = len(self.data)//10
        
        self.data["ds"] = pd.to_datetime(self.data[date_colm])
        # self.data.drop(date_colm, axis=1, inplace=True)
        self.data = self.data.astype({self.target_colm: float})

    
    def _iter_split(self, current_window: int, step_size: int):
        window_data = self.data[:current_window]
        
        if current_window + step_size > len(self.data):
            step_size = len(self.data) - current_window
        
        return window_data, step_size


    def predict(self, intial_window_size: int = None, step: int = None, freq: str = "D"):        
        window_size = intial_window_size or self.default_window_size
        step_size = step or self.default_step_size
        
        # Run iterations and return a pd series of predictions
        self.data["unique_id"] = 0
        window = window_size
        predictions = pd.Series()
        
        while window < len(self.data):
            logging.info(f"Predicting for window size: {window}")
            current_window, step_size = self._iter_split(window, step_size)
            batch_pred = self.tfm.forecast_on_df(current_window, freq=freq, value_name=self.target_colm)['timesfm']
            predictions = pd.concat([predictions, batch_pred])
            window += step_size
        supp = len(predictions) - (window - window_size)
        predictions = predictions[:-supp]
        predictions.index = [i for i in range(window_size, window)]
        return predictions