"""
pitchpredict/src/similarity
Created by Addison Kline (akline@baseball-analytica.com) in October 2024
"""
import pandas as pd
import numpy as np
from src.pitchpredict.logger_config import get_logger
from src.pitchpredict.models.pitch import Pitch

logger = get_logger(__name__)

def calculate_similarity(pitches: pd.DataFrame, context: Pitch) -> np.array:
    """
    Given a list of previous pitches, create a Similarity column that contains information on each pitch's contextual similarity to the context given.

    Args:
        pitches (DataFrame): The list of all previous pitches to add a similarity column to.
        context (Pitch): The desired context to calculate similarity to.

    Returns:
        array: The similarity column generated by this function.
    """
    logger.info('Attempting to calculate similarity scores for given pitches and context')

    batter_similarity = np.where(pitches['batter'] == context.batter_id, 1, 0)
    ball_similarity = np.where(pitches['balls'] == context.balls, 1, 0)
    strike_similarity = np.where(pitches['strikes'] == context.strikes, 1, 0)
    count_similarity = np.where((ball_similarity == 1) & (strike_similarity == 1), 1, 0)

    score_net = pitches['fld_score'] - pitches['bat_score']
    score_net_context = context.score_fld - context.score_bat
    score_similarity = 1 - 0.025*np.abs(score_net - score_net_context)

    year_similarity = 1 - 0.167*np.abs(pitches['game_year'] - context.game_year)

    final_similarity = 0.5*count_similarity + 0.3*score_similarity + 0.1*batter_similarity + 0.1*year_similarity

    logger.info('Similarity scores calculated successfully')
    return final_similarity