%% HYBRID MODEL: EV CHARGING SCHEDULE OPTIMIZATION BASED ON SOC SIMULATION
% FINAL VERSION (BUGS FIXED AND DATA UPDATED)

clc;
clear;
close all;

% Declare global variables for the CostFunction to access
global T price_grid base_load clusters delta_t num_clusters

%% ========================================================================
%  LAYER 1: MICRO-SIMULATION - TRIP AND STATE OF CHARGE (SoC)
% =========================================================================
fprintf('Running Layer 1: Micro-simulation (VinFast data)...\n');

% ---- Parameters for micro-simulation (UPDATED FROM VINFAST DATA) ----
num_sample_evs = 1000;       % Number of sample EVs for simulation
avg_battery_capacity = 73.89;   % UPDATE: Avg battery capacity (from VinFast file)
avg_efficiency = 0.16;        % UPDATE: Avg efficiency (from VinFast file)

% ---- Simulation for Home Charging Cluster (Residential) ----
% FIX: Replace normrnd with randn
dist_residential = 40 + 10 * randn([num_sample_evs, 1]); % Avg 40km, std 10
dist_residential(dist_residential <= 0) = 10; % Ensure positive distance
energy_consumed_residential = dist_residential * avg_efficiency;
initial_soc_full = 0.9;
soc_arrival_home = initial_soc_full - (energy_consumed_residential / avg_battery_capacity);
soc_arrival_home(soc_arrival_home < 0.1) = 0.1; % Assume SoC not below 10%
energy_needed_per_ev_residential = (initial_soc_full - soc_arrival_home) * avg_battery_capacity;

% ---- Simulation for Workplace Charging Cluster ----
% FIX: Replace normrnd with randn
dist_workplace = 20 + 5 * randn([num_sample_evs, 1]); % Avg 20km, std 5
dist_workplace(dist_workplace <= 0) = 5;
energy_consumed_workplace = dist_workplace * avg_efficiency;
soc_arrival_work = initial_soc_full - (energy_consumed_workplace / avg_battery_capacity);
energy_needed_per_ev_workplace = (initial_soc_full - soc_arrival_work) * avg_battery_capacity;

% Calculate average energy demand from simulation
avg_E_demand_residential = mean(energy_needed_per_ev_residential);
avg_E_demand_workplace = mean(energy_needed_per_ev_workplace);
fprintf('Micro-simulation results:\n');
fprintf('  - Average energy demand (Home Charging): %.2f kWh/EV\n', avg_E_demand_residential);
fprintf('  - Average energy demand (Workplace Charging): %.2f kWh/EV\n', avg_E_demand_workplace);

%% ========================================================================
%  LAYER 2: MACRO-OPTIMIZATION PROBLEM
% =========================================================================
fprintf('\nStarting Layer 2: Macro-optimization...\n');

%% 2.1. PROBLEM DEFINITION
T = 24;
delta_t = 1;

% 5-level Time-of-Use (ToU) price (example)
price_grid_original = [0.10,0.10,0.10,0.10,0.15,0.15,0.15,0.22,0.28,0.28,0.28,0.22,0.15,0.15,0.15,0.15,0.22,0.22,0.35,0.35,0.22,0.22,0.10,0.10]';

% ** FIX FOR CHARGING AT PEAK HOURS (18:00 & 19:00) **
% Increase electricity price at hour 18 (index 19) and 19 (index 20) significantly
% to "force" the algorithm to avoid charging at this time.
price_grid = price_grid_original;
price_grid(19) = 10.0; % 18:00 (6 PM)
price_grid(20) = 10.0; % 19:00 (7 PM)
fprintf('NOTE: Price increased for 18:00-19:00 to avoid peak charging.\n');

% Base load (example, scaled up)
base_load_small = [100, 95, 85, 85, 90, 100, 130, 180, 220, 190, 185, 180, 170, 165, 165, 175, 200, 260, 240, 200, 160, 140, 120, 105]';
scaling_factor = 10000;
base_load = base_load_small * scaling_factor;

total_num_evs = 2200000;
num_clusters = 2;
percent_residential = 0.75;
percent_workplace = 0.25;

% UPDATE: Define specific charging power based on VinFast charger file
avg_P_rated_per_ev_home = 7;      % (kW) - Assume AC home charging
avg_P_rated_per_ev_workplace = 11;  % (kW) - 11kW AC charging at workplace (from file)

% Cluster 1: Home Charging (USING RESULTS FROM LAYER 1)
num_evs_cluster1 = total_num_evs * percent_residential;
clusters(1).name = sprintf('Home Charging Group (%.2f million EVs)', num_evs_cluster1/1e6);
clusters(1).P_rated_total = num_evs_cluster1 * avg_P_rated_per_ev_home; % <-- Use P_home
clusters(1).E_demand_total = num_evs_cluster1 * avg_E_demand_residential; % <-- UPDATED FROM LAYER 1
clusters(1).T_arrival = 18; % 6 PM
clusters(1).T_depart = 8;  % 8 AM (next day)

% Cluster 2: Workplace Charging (USING RESULTS FROM LAYER 1)
num_evs_cluster2 = total_num_evs * percent_workplace;
clusters(2).name = sprintf('Workplace Charging Group (%.2f million EVs)', num_evs_cluster2/1e6);
clusters(2).P_rated_total = num_evs_cluster2 * avg_P_rated_per_ev_workplace; % <-- Use P_workplace
clusters(2).E_demand_total = num_evs_cluster2 * avg_E_demand_workplace; % <-- UPDATED FROM LAYER 1
clusters(2).T_arrival = 9;  % 9 AM
clusters(2).T_depart = 17; % 5 PM

%% 2.2. SETUP AND RUN PSO
% (Using PSO as a substitute for HSO)

CostFunction = @CostFunction_EV_HSO_VT; % Cost function handle
nVar = num_clusters * T; % Number of variables (2 clusters * 24 hours)

% Set Lower (lb) and Upper (ub) bounds for variables
lb = zeros(1, nVar);
ub = zeros(1, nVar);

for c = 1:num_clusters
    idx_start = (c-1)*T + 1;
    T_arrival = clusters(c).T_arrival;
    T_depart = clusters(c).T_depart;
    P_max_cluster = clusters(c).P_rated_total;
    
    if T_arrival > T_depart % Overnight charging (e.g., 6 PM -> 8 AM)
        for t = 1:T
            if t >= T_arrival || t < T_depart
                ub(idx_start + t - 1) = P_max_cluster;
            end
        end
    else % Daytime charging (e.g., 9 AM -> 5 PM)
        for t = 1:T
            if t >= T_arrival && t < T_depart
                ub(idx_start + t - 1) = P_max_cluster;
            end
        end
    end
end

% ---- START OPTIMIZATION ALGORITHM (EXAMPLE: PSO) ----
fprintf('Starting Optimization Algorithm (PSO)...\n');
% UPDATE: Enhance algorithm for better search
MaxIt = 200;     % Max iterations (increased from 100)
nPop = 80;       % Population size (increased from 50)
w = 0.729;       % Inertia weight
c1 = 1.494;      % Personal coefficient
c2 = 1.494;      % Social coefficient

% Initialization
particle.Position = [];
particle.Velocity = [];
particle.Cost = [];
particle.Best.Position = [];
particle.Best.Cost = [];
pop = repmat(particle, nPop, 1);
GlobalBest.Cost = inf;

for i = 1:nPop
    pop(i).Position = lb + (ub - lb) .* rand(1, nVar);
    pop(i).Velocity = zeros(1, nVar);
    pop(i).Cost = CostFunction(pop(i).Position);
    pop(i).Best.Position = pop(i).Position;
    pop(i).Best.Cost = pop(i).Cost;
    if pop(i).Best.Cost < GlobalBest.Cost
        GlobalBest = pop(i).Best;
    end
end

BestCost = zeros(MaxIt, 1);

% PSO Main Loop
for it = 1:MaxIt
    for i = 1:nPop
        pop(i).Velocity = w * pop(i).Velocity ...
            + c1 * rand(1, nVar) .* (pop(i).Best.Position - pop(i).Position) ...
            + c2 * rand(1, nVar) .* (GlobalBest.Position - pop(i).Position);
        pop(i).Position = pop(i).Position + pop(i).Velocity;
        pop(i).Position = max(pop(i).Position, lb);
        pop(i).Position = min(pop(i).Position, ub);
        pop(i).Cost = CostFunction(pop(i).Position);
        if pop(i).Cost < pop(i).Best.Cost
            pop(i).Best.Position = pop(i).Position;
            pop(i).Best.Cost = pop(i).Cost;
            if pop(i).Best.Cost < GlobalBest.Cost
                GlobalBest = pop(i).Best;
            end
        end
    end
    BestCost(it) = GlobalBest.Cost;
    if mod(it, 20) == 0 % Print every 20 iterations to reduce clutter
        fprintf('Iteration %d/%d, Best Cost = %e\n', it, MaxIt, BestCost(it));
    end
end
fprintf('Optimization completed.\n');
BestSolution = GlobalBest.Position;

%% 2.3. VISUALIZE RESULTS
fprintf('Processing and plotting results...\n');
% Extract optimal solution
P_ev = reshape(BestSolution, T, num_clusters)'; % Row = cluster, Col = hour
P_ev_cluster1 = P_ev(1, :);
P_ev_cluster2 = P_ev(2, :);
P_ev_total = sum(P_ev, 1)'; % Total EV charging load per hour
total_load = base_load + P_ev_total; % Total grid load

time_axis = (0:T-1)';

% ========================================================================
% PLOT ANNOTATIONS
% ========================================================================

% ---- Figure 1: STACKED Load Profile ----
% Purpose: Show how the total grid load has been "valley-filled".
% - GRAY Area: Original base load.
% - BLUE Area: Optimized EV charging load.
% - RED Line: Final total load (Gray + Blue).
% -> A good result is a flatter RED line compared to the GRAY area.
figure(1);
set(gcf, 'Name', 'Fig.1: Load Profile Optimization (Stacked)', 'NumberTitle', 'off');
hold on; grid on;
area(time_axis, base_load, 'FaceColor', [0.8 0.8 0.8], 'DisplayName', 'Base Load');
area(time_axis, base_load + P_ev_total, 'FaceColor', [0.2 0.6 1.0], 'DisplayName', 'EV Load (Optimized)');
plot(time_axis, total_load, 'r-', 'LineWidth', 2.5, 'DisplayName', 'Total Load');
title('Grid Load Before and After EV Charging Optimization');
xlabel('Hour of Day');
ylabel('Power (kW)');
legend('show', 'Location', 'northwest');
% FIX: Change 23 to 24 to show full 24 hours
set(gca, 'XTick', 0:2:24);
xlim([0 24]);
hold off;

% ---- Figure 2: Detailed Charging Schedule by CLUSTER ----
% Purpose: Show how the algorithm allocated charging for the 2 different groups.
% - Top Plot (Home): Charges at night (low-price hours, avoided 18-19h).
% - Bottom Plot (Workplace): Charges at midday (low-price hours 12-15h).
figure(2);
set(gcf, 'Name', 'Fig.2: Optimal Charging Schedule by Cluster', 'NumberTitle', 'off');
subplot(2,1,1);
bar(time_axis, P_ev_cluster1, 'b');
title(clusters(1).name);
xlabel('Hour of Day'); ylabel('Charging Power (kW)');
% FIX: Change 23 to 24
set(gca, 'XTick', 0:2:24); xlim([0 24]);
subplot(2,1,2);
bar(time_axis, P_ev_cluster2, 'r');
title(clusters(2).name);
xlabel('Hour of Day'); ylabel('Charging Power (kW)');
% FIX: Change 23 to 24
set(gca, 'XTick', 0:2:24); xlim([0 24]);

% ---- Figure 3: Correlation between PRICE and TOTAL CHARGING ----
% Purpose: Check if the algorithm charges during low-price hours.
% - RED Line (left-axis): Electricity price (fixed, very high at 18-19h).
% - BLUE Bars (right-axis): Total EV charging power.
% -> Good result: BLUE bars only appear when RED line is low.
figure(3);
set(gcf, 'Name', 'Fig.3: Price vs. Charging Schedule', 'NumberTitle', 'off');
yyaxis left;
% Use price_grid_original to plot (excluding the 10.0 spikes) for clarity
plot(time_axis, price_grid_original, 'r-o', 'LineWidth', 2, 'DisplayName', 'Electricity Price (Original)');
ylabel('Price (USD/kWh)');
ylim([0 max(price_grid_original)*1.2]); % Set y-axis based on original price
yyaxis right;
bar(time_axis, P_ev_total, 'b', 'DisplayName', 'EV Load');
ylabel('Charging Power (kW)');
title('EV Charging Schedule vs. Electricity Price');
legend('show', 'Location', 'north');
% FIX: Change 23 to 24
set(gca, 'XTick', 0:2:24); xlim([0 24]);

% ---- Figure 4: Algorithm Convergence Plot ----
% Purpose: Show the algorithm's learning process finding the best cost.
% - LINE: Cost decreases with each iteration.
% -> Good result: The line flattens at the end, showing
%    the algorithm has found the best possible optimum.
figure(4);
set(gcf, 'Name', 'Fig.4: Convergence Plot', 'NumberTitle', 'off');
plot(1:MaxIt, BestCost, 'k-', 'LineWidth', 2);
title('Algorithm Convergence');
xlabel('Iteration');
ylabel('Total Cost (USD)');
grid on;

% ---- Figure 5: SEPARATED Load Profiles (As requested) ----
% Purpose: Split Figure 1 into 2 parts for easier analysis.
% - Top Plot: Only base load (GRAY area).
% - Bottom Plot: Only EV load (BLUE area).
% -> Clearly shows the BLUE area was created to "fill" the valley of the GRAY area.
figure(5); 
set(gcf, 'Name', 'Fig.5: Separated Load Components', 'NumberTitle', 'off');
subplot(2,1,1);
bar(time_axis, base_load, 'FaceColor', [0.8 0.8 0.8]);
title('Component 1: Base Load');
xlabel('Hour of Day');
ylabel('Power (kW)');
% FIX: Change 23 to 24
set(gca, 'XTick', 0:2:24);
xlim([0 24]);
grid on;
subplot(2,1,2);
bar(time_axis, P_ev_total, 'FaceColor', [0.2 0.6 1.0]);
title('Component 2: EV Load (Optimized)');
xlabel('Hour of Day');
ylabel('Power (kW)');
% FIX: Change 23 to 24
set(gca, 'XTick', 0:2:24);
xlim([0 24]);
grid on;

% ---- Figure 7: Results from LAYER 1 (Micro-simulation) ----
% (No Figure 6)
% Purpose: Display the results of the trip simulation.
% - Top Plot: SoC (State of Charge) distribution upon arrival at charger.
% - Bottom Plot: Energy demand (kWh) distribution for each EV.
figure(7);
clf;
set(gcf, 'Name', 'Fig.7: Micro-Simulation Results', 'NumberTitle', 'off');
% Subplot 1: SoC distribution on arrival
subplot(2, 1, 1);
hold on; grid on;
histogram(soc_arrival_home, 20, 'FaceColor', 'b', 'Normalization', 'pdf');
histogram(soc_arrival_work, 20, 'FaceColor', 'r', 'Normalization', 'pdf');
title('State of Charge (SoC) Distribution on Arrival');
xlabel('SoC'); ylabel('Probability Density');
legend('Arriving Home', 'Arriving Workplace');
hold off;
% Subplot 2: Energy demand distribution
subplot(2, 1, 2);
hold on; grid on;
histogram(energy_needed_per_ev_residential, 20, 'FaceColor', 'b', 'Normalization', 'pdf');
histogram(energy_needed_per_ev_workplace, 20, 'FaceColor', 'r', 'Normalization', 'pdf');
title('Charging Energy Demand Distribution per EV');
xlabel('Energy Needed (kWh)'); ylabel('Probability Density');
legend('Home Charging', 'Workplace Charging');
hold off;

fprintf('Done! Generated 6 plot windows (Fig 1, 2, 3, 4, 5, 7).\n');