import numpy as np
import matplotlib.pyplot as plt


class NadarayaWatson():
    def __init__() -> None: pass

    @staticmethod
    def gauss_kernel(u):
        return np.exp(-0.5 * u**2) / np.sqrt(2 * np.pi)

    @staticmethod
    def nadaraya_watson_envelope_indicator(prices,
                                           bandwidth: int = 10,
                                           env_percent: int = 1):
        n = len(prices)
        envelopes = np.zeros((n, 3))  # Two columns for upper and lower envelopes

        for i in range(n):
            # Estimate conditional expectation using Nadaraya-Watson estimator
            numerator = \
                np.sum(NadarayaWatson.gauss_kernel(
                    (i - np.arange(n)) / bandwidth) * prices)
            denominator = \
                np.sum(NadarayaWatson.gauss_kernel(
                    (i - np.arange(n)) / bandwidth))
            conditional_expectation = \
                numerator / denominator if denominator != 0 else 0
            
            # Calculate upper and lower envelopes
            envelopes[i, 0] = \
                conditional_expectation + np.std(prices) * env_percent
            envelopes[i, 1] = \
                conditional_expectation - np.std(prices) * env_percent
            envelopes[i, 2] = \
                conditional_expectation
        return envelopes

    # @staticmethod
    # def get_viz():
    #     # Plotting
    #     plt.figure(figsize=(10, 6))
    #     plt.plot(self.prices, label='Prices', color='blue')
    #     plt.plot(self.envelopes[:, 0], label='Upper Envelope', color='red', linestyle='--')
    #     plt.plot(self.envelopes[:, 1], label='Lower Envelope', color='green', linestyle='--')
    #     plt.plot(self.envelopes[:, 2], label='Estimator', color='k', linestyle='--')
    #     plt.title('Nadaraya-Watson Envelope Indicator')
    #     plt.xlabel('Days')
    #     plt.ylabel('Price')
    #     plt.legend()
    #     plt.grid(True)
    #     plt.show()

# Example Usage 
# if __name__ == "__main__":
#     data = vns.stock_historical_data(symbol="ACB", start_date="2022-01-31", 
#                                 end_date='2024-01-31', resolution='1D',
#                                 type='stock', beautify=True, decor=False, 
#                                 source='DNSE').iloc[-269:-68].reset_index()
#     nada = NadarayaWatson(data)
#     print(nada.envelopes[-1])



# %%