import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import itertools
from sklearn.linear_model import LinearRegression

def convert_to_am_pm(time_str):
    try:
        # Check if the input is a string
        if not isinstance(time_str, str):
            raise ValueError("Input must be a string.")

        # Remove any leading/trailing whitespace and convert to lowercase
        time_str = time_str.strip().lower()

        # Check for 'am' or 'pm' in the input
        if 'am' in time_str or 'pm' in time_str:
            period = 'AM' if 'am' in time_str else 'PM'
            time_str = time_str.replace('am', '').replace('pm', '').strip()
            if ':' not in time_str:
                hours = int(time_str)
                minutes = 0
            else:
                parts = time_str.split(":")
                hours = int(parts[0])
                minutes = int(parts[1])
        else:
            # If no 'am' or 'pm', assume 24-hour format
            if ':' not in time_str:
                hours = int(time_str)
                minutes = 0
            else:
                parts = time_str.split(":")
                hours = int(parts[0])
                minutes = int(parts[1])
            period = 'AM' if hours < 12 else 'PM'
            if hours == 0:
                hours = 12
            elif hours > 12:
                hours -= 12

        # Validate hours and minutes
        if hours < 1 or hours > 12:
            raise ValueError("Hours must be between 1 and 12.")
        if minutes < 0 or minutes > 59:
            raise ValueError("Minutes must be between 0 and 59.")

        # Format the time
        formatted_time = f"{hours:02}:{minutes:02} {period}"
        return formatted_time

    except ValueError as e:
        return str(e)

# Function to convert time to minutes since new start
def convert_time_to_minutes_from_zero(time_list, zero_time):
    return [(datetime.strptime(t, '%I:%M %p') - zero_time).total_seconds() / 60 for t in time_list]

# Function to plot updated datasets with a combined trend line and horizontal
# line at 100%
def plot_updated_datasets_with_combined_trendline(datasets):
    plt.figure(figsize=(12, 8))
    colors = itertools.cycle(['blue', 'orange', 'green', 'red', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan'])
    combined_time = []
    combined_percentage = []
    for label, data in datasets.items():
        original_start_time = datetime.strptime(data['Time'][0], '%I:%M %p')
        time_minutes = convert_time_to_minutes_from_zero(data['Time'], original_start_time)
        df = pd.DataFrame({'Time': time_minutes, 'Percentage': data['Percentage']})
        # Linear regression for individual dataset
        X_individual = np.array(time_minutes).reshape(-1, 1)
        y_individual = np.array(data['Percentage'])
        model_individual = LinearRegression().fit(X_individual, y_individual)
        # Predict the time to reach 0% for individual dataset
        time_to_0 = -model_individual.intercept_ / model_individual.coef_[0]
        zero_time = original_start_time + timedelta(minutes=time_to_0)
        # Add the predicted zero time to the dataset at the beginning
        data['Time'] = [zero_time.strftime('%I:%M %p')] + data['Time']
        data['Percentage'] = [0] + data['Percentage']
        # Update dataframe with new data
        time_minutes_updated = convert_time_to_minutes_from_zero(data['Time'], zero_time)
        df_updated = pd.DataFrame({'Time': time_minutes_updated, 'Percentage': data['Percentage']})
        plt.plot(df_updated['Time'], df_updated['Percentage'], 'o-', label=label, color=next(colors))
        # Append to combined data for trend line
        combined_time.extend(time_minutes_updated)
        combined_percentage.extend(data['Percentage'])

    # Combined trend line
    X_combined = np.array(combined_time).reshape(-1, 1)
    y_combined = np.array(combined_percentage)
    model_combined = LinearRegression().fit(X_combined, y_combined)
    trendline_combined = model_combined.predict(X_combined)
    plt.plot(combined_time, trendline_combined, "r--", label='Combined Trendline')
    # Predict the time to reach 100% using combined data
    time_to_100_combined = (100 - model_combined.intercept_) / model_combined.coef_[0]
    predicted_time_to_100_combined = zero_time + timedelta(minutes=time_to_100_combined)
    # Add horizontal line at 100% and annotate
    plt.axhline(y=100, color='green', linestyle='--', label='100% Line')
    plt.axvline(x=time_to_100_combined, color='green', linestyle='--')
    # plt.text(time_to_100_combined, 50, f'100% at {predicted_time_to_100_combined.strftime("%I:%M %p")}', color='green')
    plt.xlabel('Time (minutes since 0% start)')
    plt.ylabel('Battery Percentage')
    plt.title('Updated Battery Percentage Over Time with Combined Trendline')
    plt.legend()
    plt.grid(True)
    plt.show()
    return model_combined, zero_time

# Function to predict the time to reach 100% battery based on user input
def predict_battery_time(model, p1, t1_str, zero_time):
    current_time = datetime.strptime(t1_str, '%I:%M %p')
    time_since_zero = (current_time - zero_time).total_seconds() / 60
    time_to_100 = (100 - p1) / model.coef_[0]
    total_time_to_100 = time_since_zero + time_to_100
    predicted_time_to_100 = zero_time + timedelta(minutes=total_time_to_100)
    return total_time_to_100, predicted_time_to_100.strftime('%I:%M %p')

# Function to predict battery for the given time
def predict_battery_percentage_at_time(model, p1, t1_str, zero_time, t2_str):
    t2 = datetime.strptime(t2_str, '%I:%M %p')
    time_since_zero_t2 = (t2 - zero_time).total_seconds() / 60
    time_since_t1 = (t2 - datetime.strptime(t1_str, '%I:%M %p')).total_seconds() / 60
    p2 = p1 + model.coef_[0] * time_since_t1
    return p2

# Example datasets
datasets = {
    'Dataset 1': {
        'Time': ['09:00 AM', '09:33 AM', '10:48 AM', '12:01 PM', '12:31 PM', '12:38 PM', '01:12 PM', '01:27 PM', '01:42 PM', '01:50 PM', '02:00 PM', '02:23 PM'],
        'Percentage': [28, 33, 50, 68, 74, 76, 83, 87, 90, 91, 93, 97]
    },
    'Dataset 2': {
        'Time': ['09:39 AM', '09:55 AM', '10:22 AM', '10:48 AM', '11:11 AM', '11:23 AM', '12:04 PM', '12:11 PM', '12:40 PM', '1:18 PM', '1:37 PM', '1:55 PM', '2:01 PM', '2:28 PM', '2:53 PM'],
        'Percentage': [21, 24, 29, 36, 41, 44, 54, 55, 63, 70, 75, 79, 80, 86, 91]
    }
}

# Plot the updated datasets with a combined trend line
model_combined, zero_time = plot_updated_datasets_with_combined_trendline(datasets)

# Take user input
p1 = float(input("Enter the current battery percentage (p1): "))
t1_str = convert_to_am_pm(input("Enter the current time (t1) in 'hh:mm AM/PM' format: "))
t2_str = convert_to_am_pm(input("Enter the prediction time (t2) in 'hh:mm AM/PM' format: "))
print("Considering times as " + t1_str + " and " + t2_str)
# print ("Considering time as " + t1_str)

# Predict the time to reach 100% battery
total_time_to_100, predicted_time_to_100 = predict_battery_time(model_combined, p1, t1_str, zero_time)
print(f"The battery will reach 100% at {predicted_time_to_100}, which is {total_time_to_100:.2f} minutes from the zero start time.")

# Predict the battery % for time t2
predicted_percentage_t2 = predict_battery_percentage_at_time(model_combined, p1, t1_str, zero_time, t2_str)
print(f"The battery percentage at {t2_str} is predicted to be {predicted_percentage_t2:.2f}%.")