import pandas as pd
import numpy as np
import geopandas as gpd
import shap
import xgboost as xgb
import folium
from folium.plugins import HeatMap
import branca.colormap as cm
import matplotlib.pyplot as plt
from shapely.geometry import Point
from sklearn.model_selection import train_test_split

# --- Step 1: Load Dataset ---
# URL for the California Housing dataset
url = "https://raw.githubusercontent.com/ageron/handson-ml/master/datasets/housing/housing.csv"
print("Status: Downloading dataset...")
df = pd.read_csv(url)

print(f"Success: Dataset loaded with {len(df)} records.")
# Display first 5 rows
df.head()

# GIS Feature Engineering
# Calculate the precise geodesic distance from each property to the nearest coastline.

# Convert DataFrame to GeoDataFrame
# Initial CRS: EPSG:4326 (WGS84 - Latitude/Longitude)
gdf = gpd.GeoDataFrame(
    df, geometry=gpd.points_from_xy(df.longitude, df.latitude), crs="EPSG:4326"
)

print("Fetching California coastline data...")
# Load California counties boundary from public GeoJSON
ca_url = "https://raw.githubusercontent.com/codeforgermany/click_that_hood/main/public/data/california-counties.geojson"
ca_counties = gpd.read_file(ca_url)

# Topology Repair
print("Repairing topology errors (Buffer 0 fix)...")
# Fix self-intersection or invalid geometries to prevent dissolve errors
ca_counties['geometry'] = ca_counties.geometry.buffer(0)

# Dissolve counties to get the outer boundary of California
print("Dissolving boundaries...")
ca_boundary = ca_counties.dissolve()

# Reproject to EPSG:3310 (California Albers) for accurate meter-based distance calculation
print("Reprojecting to EPSG:3310 (Meters)...")
gdf = gdf.to_crs("EPSG:3310")
ca_boundary = ca_boundary.to_crs("EPSG:3310")

# Extract the boundary line
coastline = ca_boundary.boundary

print("Calculating distance to coast")
# Calculate distance for every point
gdf['dist_to_coast'] = gdf.geometry.apply(lambda x: coastline.distance(x))

# Convert back to standard DataFrame for Machine Learning
df_ml = pd.DataFrame(gdf.drop(columns='geometry'))
# Convert meters to kilometers for better readability
df_ml['dist_to_coast_km'] = df_ml['dist_to_coast'] / 1000 

print("GIS Feature Engineering complete.")
print(df_ml[['longitude', 'latitude', 'median_house_value', 'dist_to_coast_km']].head())

# Data Preparation & Modeling 

# Feature Selection:
# Remove target variable, old categorical column, and the raw meter-distance column
X = df_ml.drop(columns=['median_house_value', 'ocean_proximity', 'dist_to_coast'])

# Handle missing values (Simple imputation with mean)
X = X.fillna(X.mean())

y = df_ml['median_house_value']

# Split data into Training and Testing sets (80/20 split)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Initialise and Train XGBoost Regressor
print("Training XGBoost Model...")
model = xgb.XGBRegressor(n_estimators=100, learning_rate=0.1, max_depth=6, random_state=42)
model.fit(X_train, y_train)

# Evaluate Model
score = model.score(X_test, y_test)
print(f"Success: Model R² Score: {score:.4f}")

# Model Explainability (SHAP Analysis)

print("Status: Calculating SHAP values...")

# Initialise TreeExplainer
explainer = shap.TreeExplainer(model)
shap_values = explainer.shap_values(X_test)

# Visualisation: Optimized Beeswarm Plot
# Downsampling to avoid overplotting and improve clarity
limit = 2000
if X_test.shape[0] > limit:
    indices = np.random.choice(X_test.shape[0], limit, replace=False)
    shap_subset = shap_values[indices]
    X_subset = X_test.iloc[indices]
else:
    shap_subset = shap_values
    X_subset = X_test

# Plotting
plt.figure(figsize=(12, 8))
shap.summary_plot(shap_subset, X_subset, 
                  max_display=10, 
                  alpha=0.6, 
                  plot_size=(10, 6),
                  show=False)

# Add Labels
plt.title("Feature Impact on House Prices", fontsize=16, pad=20)
plt.xlabel("SHAP Value (Impact on Price Output)", fontsize=12)
plt.yticks(fontsize=12)

print("Displaying SHAP Summary Plot:")
plt.show()

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

shap.summary_plot(shap_values, X_test, 
                  plot_type="bar", 
                  max_display=10, 
                  color='#007acc', 
                  show=False)

plt.title("Average Impact Magnitude of Features", fontsize=16, pad=20)
plt.xlabel("Average |SHAP Value| (Mean Impact on Price)", fontsize=12)
plt.show()

# Interactive Geospatial Visualization
# Visualise the "Location Premium" (Pure spatial contribution to price)

print("Generating interactive map with legend...")

# 1. Sampling for performance
sample_size = 2000 
actual_size = min(sample_size, len(X_test))
indices = np.random.choice(X_test.index, size=actual_size, replace=False)

# 2. Extract Data
lat_idx = X_test.columns.get_loc('latitude')
lon_idx = X_test.columns.get_loc('longitude')
dist_idx = X_test.columns.get_loc('dist_to_coast_km')

subset_shap = shap_values[np.isin(X_test.index, indices)]

# Calculate Location Premium: Sum of Latitude, Longitude, and Distance-to-Coast effects
location_impact = subset_shap[:, lat_idx] + subset_shap[:, lon_idx] + subset_shap[:, dist_idx]

# Define range for the legend
min_val = np.min(location_impact)
max_val = np.max(location_impact)

# Prepare coordinates
lat_data = X_test.loc[indices, 'latitude'].values
lon_data = X_test.loc[indices, 'longitude'].values

# 3. Create Base Map
m = folium.Map(location=[36.7, -119.4], zoom_start=6, tiles='CartoDB dark_matter')

# 4. Add Color Legend (Red = High Premium, Blue = Negative Impact)
colormap = cm.LinearColormap(
    colors=['blue', 'cyan', 'lime', 'yellow', 'red'],
    vmin=min_val,
    vmax=max_val,
    caption='Location Premium (Net Impact on Price in USD)' 
)
m.add_child(colormap)

# 5. Prepare Data for HeatMap 
data_heat = []
for i in range(len(lat_data)):
    # Explicitly convert numpy types to standard python floats
    lat = float(lat_data[i])
    lon = float(lon_data[i])
    weight = float(location_impact[i])
    
    # Normalize weight to 0-1 range for the heatmap gradient
    normalized_weight = (weight - min_val) / (max_val - min_val)
    
    data_heat.append([lat, lon, normalized_weight])

# 6. Add HeatMap Layer
HeatMap(data_heat, radius=15, blur=20, 
        gradient={0.0: 'blue', 0.25: 'cyan', 0.5: 'lime', 0.75: 'yellow', 1.0: 'red'}).add_to(m)

print(f"Map generated. Value range: ${min_val:.0f} to ${max_val:.0f}")
m