import os
import pandas as pd
import numpy as np
import datetime
from datetime import datetime, date
import warnings
warnings.filterwarnings("ignore")

import snowflake.connector
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization

from password import PRIVATE_KEY_PASSPHRASE

def fix_time_variables(df):
    df['FIXED_RELEASE_DATE'] = pd.to_datetime(df['FIXED_RELEASE_DATE'])
    df['MIN_CHART_DATE'] = pd.to_datetime(df['MIN_CHART_DATE'])
    df['LATEST_HIT_DATE'] = pd.to_datetime(df['LATEST_HIT_DATE'])
    df['Month_release'] = df['FIXED_RELEASE_DATE'].dt.month_name()
    df['Month_chart'] = df['MIN_CHART_DATE'].dt.month_name()
    df['Month_latest_chart'] = df['LATEST_HIT_DATE'].dt.month_name()
    return df

def prepare_sample(df):
    prefixes = ['USS1Z', 'QMS1Z', 'ISRC', 'QZS1Z',  'JMK40', 'AB'] 
    mask = df['ISRC'].str.startswith(tuple(prefixes))
    filtered = df[~mask]
    filtered['FIXED_RELEASE_DATE'] = filtered['FIXED_RELEASE_DATE'].dt.date
    sub = filtered[filtered['FIXED_RELEASE_DATE'] <= date(2024,2, 13)]
    return sub

def artist_score_calculation(df):
    conditions = [(df['PREVIOUS_HIT']==0),
                 (df['PREVIOUS_HIT']==1) & (df['DAYS_SINCE_LAST_HIT'] > 365),
                (df['PREVIOUS_HIT']==1) & (df['DAYS_SINCE_LAST_HIT'] <= 365)]
    choices = ['no_previous_hit', 'non_recent_hit', 'recent_hit']
    df['artist_score'] = np.select(conditions, choices)
    return df

def fix_dataframe(df):
    df = artist_score_calculation(df)
    sub = prepare_sample(df)
    return sub

def prepare_data_for_regression(df):
    country_transformed = df.copy()
    scaled = []
    for scaler in pickled_scalers:
        features = scaler.feature_names_in_
        scaled = scaled + [pd.DataFrame(scaler.transform(country_transformed[features]), columns = features)]
    country_transformed = pd.concat(scaled, axis=1)
    return country_transformed

def calculations_country(country_transformed, country):
    df = fix_time_variables(df)
    df = artist_score_calculation(df)
    country_transformed = prepare_data_for_regression(df)
    predicted = pickled_model.predict(country_transformed)
    probs = pickled_model.predicted_probabilities(country_transformed)
    COUNTRY = country
    res = pd.DataFrame({'Run_date': datetime.now(),
                    'COUNTRY_SAMPLE': COUNTRY,
                    'ISRC': df['ISRC'], 
                    'Artist_name': df['ARTIST_NAME'],
                    'Release_date':df['FIXED_RELEASE_DATE'],
                    'Predicted_probability_entry': probs[:,1],
                    'Prediction': predicted,
                    'Latest_chart_date': df['MAX_CHART_DATE'],
                    'Chart_country': df['CHART_COUNTRY'],
                    'previous_hit_in_country_of_interest': df['artist_score'],
                    'ACOUSTICNESS': df['ACOUSTICNESS'], 
                    'DANCEABILITY': df['DANCEABILITY'], 
                    'DURATION_MS': df['DURATION_MS'],
                    'ENERGY': df['ENERGY'], 
                    'INSTRUMENTALNESS': df['INSTRUMENTALNESS'], 
                    'KEY': df['KEY'], 
                    'LIVENESS': df['LIVENESS'], 
                    'LOUDNESS': df['LOUDNESS'], 
                    'MODE': df['MODE'],
                    'SPEECHINESS': df['SPEECHINESS'], 
                    'TEMPO': df['TEMPO'], 
                    'TIME_SIGNATURE': df['TIME_SIGNATURE'],
                    'VALENCE': df['VALENCE']})
    return res 
    

if __name__ == '__main__':

    with open("/Users/impr001/Keys_snowflake/rsa_key.p8", "rb") as key: 
        p_key= serialization.load_pem_private_key(
            key.read(),
            password=PRIVATE_KEY_PASSPHRASE.encode(),
            backend=default_backend()
        )

    pkb = p_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption())

    print("We reach that point here ...")
    ctx = snowflake.connector.connect(
        user='eimpara@theorchard.com', 
        account='orchard',
        private_key=pkb,
        role= 'PROD_DATALYTICS_ROLE',
        warehouse = "DEV_ANALYTICS_ORCHARD"
        )
    print("And here...")

    try:
        cs = ctx.cursor()
        sql = """
        select * from intelligence.dbt_prod.orch_spotify_top_200_chart_appearances
        where ISRC = 'ITPT82400001';
        """
        cs.execute(sql)
        print("executed")
        df = cs.fetch_pandas_all()

        print(df.shape)

    finally:
        cs.close()

    with open('/Users/impr001/Documents/Jupyter_Notebooks/Maze/log_model_italy_v3.pkl', 'rb') as fileIT:
       (pickled_model, pickled_scalers) = pd.compat.pickle_compat.load(fileIT)
        
    res = calculations_country(df, 'Italy')
    print(res.head())