import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error

# ============================================
# Part 1: Load and Prepare Dataset
# ============================================

# Given dataset
data = {
    "Storage": [32, 32, 64, 64, 128, 128, 128, 256, 256, 256, 512, 512, 512, 512, 512],
    "Price": [10, 12, 15, 18, 22, 25, 30, 35, 38, 42, 50, 55, 60, 65, 70]
}

# Convert to DataFrame
df_phones = pd.DataFrame(data)

print("="*60)
print("Q2: Smartphone Price Prediction - Linear Regression")
print("="*60)

print("\n Dataset Overview:")
print(df_phones.head(10))
print(f"\n Dataset Shape: {df_phones.shape}")
print(f" Storage Range: {df_phones['Storage'].min()}GB - {df_phones['Storage'].max()}GB")
print(f" Price Range: ₹{df_phones['Price'].min()}K - ₹{df_phones['Price'].max()}K")

# Check for any missing values
print(f"\n Missing Values:\n{df_phones.isnull().sum()}")

# ============================================
# Part 2: Split Features and Target
# ============================================

# Independent variable (X) - Storage
X = df_phones[['Storage']]  # 2D array (required by sklearn)

# Dependent variable (y) - Price
y = df_phones['Price']      # 1D array

print("\n Feature (Independent Variable): Storage")
print(f"Shape: {X.shape}")

print("\n Target (Dependent Variable): Price")
print(f"Shape: {y.shape}")

# ============================================
# Part 3: Split into Training and Testing Sets
# ============================================

# Split data: 80% training, 20% testing
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

print(f"\n Training Set Size: {len(X_train)} samples")
print(f" Testing Set Size: {len(X_test)} samples")

# ============================================
# Part 4: Build and Train Linear Regression Model
# ============================================

# Create and train the model
model = LinearRegression()
model.fit(X_train, y_train)

# Get model coefficients
print(f"\n Model Coefficients:")
print(f"Intercept (β0): {model.intercept_:.4f}")
print(f"Coefficient (β1): {model.coef_[0]:.4f}")
print(f"\n Equation: Price = {model.intercept_:.2f} + {model.coef_[0]:.2f} * Storage")

# ============================================
# Part 5: Predict Prices
# ============================================

# Predict on training and test data
y_train_pred = model.predict(X_train)
y_test_pred = model.predict(X_test)

# Predict for new storage values (Indian market scenario)
new_storage = np.array([[64], [128], [256], [512], [1024]])
predicted_prices = model.predict(new_storage)

print("\n Predictions for New Storage Options:")
predictions_df = pd.DataFrame({
    'Storage (GB)': new_storage.flatten(),
    'Predicted Price (₹K)': np.round(predicted_prices, 2)
})
print(predictions_df)

# ============================================
# Part 6: Visualize Regression Line on Scatter Plot
# ============================================

plt.figure(figsize=(12, 6))

# Scatter plot of actual data
plt.scatter(X, y, color='blue', s=100, alpha=0.6, label='Actual Data', edgecolors='black')

# Regression line
x_range = np.linspace(X.min(), X.max(), 100).reshape(-1, 1)
y_range = model.predict(x_range)
plt.plot(x_range, y_range, color='red', linewidth=3, label='Regression Line')

# Highlight training and test points
plt.scatter(X_train, y_train, color='green', s=80, alpha=0.5, marker='o', label='Training Data')
plt.scatter(X_test, y_test, color='orange', s=120, alpha=0.7, marker='s', label='Test Data')

# Plot predictions
plt.scatter(new_storage, predicted_prices, color='purple', s=200, marker='D', 
            label='Predicted for New Storage', edgecolors='black', zorder=5)

# Add value labels for predictions
for i, (storage, price) in enumerate(zip(new_storage.flatten(), predicted_prices)):
    plt.annotate(f'{storage}GB: ₹{price:.1f}K', 
                 (storage, price), 
                 xytext=(10, 10), 
                 textcoords='offset points',
                 fontsize=9,
                 bbox=dict(boxstyle='round,pad=0.3', facecolor='yellow', alpha=0.7))

# Customize plot
plt.xlabel('Storage (GB)', fontsize=12)
plt.ylabel('Price (₹K)', fontsize=12)
plt.title('Smartphone Price Prediction\nStorage vs Price (Linear Regression)', fontsize=14, fontweight='bold')
plt.legend(loc='upper left', fontsize=10)
plt.grid(True, alpha=0.3)
plt.xlim(0, 1100)
plt.ylim(0, 120)

# Add annotation for model equation
plt.text(550, 20, f'Price = {model.intercept_:.2f} + {model.coef_[0]:.2f} * Storage', 
         fontsize=11, bbox=dict(boxstyle='round,pad=0.5', facecolor='lightgray', alpha=0.8))

plt.show()

# ============================================
# Part 7: Evaluate Model Performance
# ============================================

print("\n" + "="*60)
print("PART 5: Model Evaluation")
print("="*60)

# R² Score
r2_train = r2_score(y_train, y_train_pred)
r2_test = r2_score(y_test, y_test_pred)

print(f"\n R² Score (Training): {r2_train:.4f}")
print(f" R² Score (Testing): {r2_test:.4f}")

# Other metrics
mse_train = mean_squared_error(y_train, y_train_pred)
mse_test = mean_squared_error(y_test, y_test_pred)
rmse_train = np.sqrt(mse_train)
rmse_test = np.sqrt(mse_test)
mae_train = mean_absolute_error(y_train, y_train_pred)
mae_test = mean_absolute_error(y_test, y_test_pred)

print(f"\n Training Metrics:")
print(f"   MSE: {mse_train:.4f}")
print(f"   RMSE: {rmse_train:.4f}")
print(f"   MAE: {mae_train:.4f}")

print(f"\n Testing Metrics:")
print(f"   MSE: {mse_test:.4f}")
print(f"   RMSE: {rmse_test:.4f}")
print(f"   MAE: {mae_test:.4f}")

# ============================================
# Part 8: Residual Analysis
# ============================================

# Calculate residuals
residuals = y_test - y_test_pred

# Plot residuals
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.scatter(y_test_pred, residuals, color='purple', alpha=0.7, s=80)
plt.axhline(y=0, color='red', linestyle='--', linewidth=2)
plt.xlabel('Predicted Price (₹K)', fontsize=12)
plt.ylabel('Residuals', fontsize=12)
plt.title('Residual Plot', fontsize=12)
plt.grid(True, alpha=0.3)

plt.subplot(1, 2, 2)
plt.hist(residuals, bins=8, color='lightgreen', edgecolor='black', alpha=0.7)
plt.xlabel('Residuals', fontsize=12)
plt.ylabel('Frequency', fontsize=12)
plt.title('Distribution of Residuals', fontsize=12)
plt.axvline(x=0, color='red', linestyle='--', linewidth=2)
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# ============================================
# Part 9: Real-world Indian Market Insights
# ============================================

print("\n" + "="*60)
print("Real-World Indian Market Insights")
print("="*60)

# Predict prices for popular Indian smartphone configurations
indian_configs = np.array([[64], [128], [256], [512]])
indian_prices = model.predict(indian_configs)

print("\n Predicted Prices for Indian Smartphone Configurations:")
config_data = pd.DataFrame({
    'Storage': indian_configs.flatten(),
    'Predicted Price (₹K)': np.round(indian_prices, 2),
    'Predicted Price (₹)': np.round(indian_prices * 1000, 0)
})
print(config_data)

print("\n Price Comparison:")
print("64GB: ₹15,000 - ₹18,000 (Budget segment)")
print("128GB: ₹22,000 - ₹30,000 (Mid-range)")
print("256GB: ₹35,000 - ₹42,000 (Premium)")
print("512GB: ₹50,000 - ₹70,000 (Flagship)")

print("\n Model Training and Evaluation Complete!")