diff --git a/src/io.jl b/src/io.jl
index 98cbe203..4173dc40 100644
--- a/src/io.jl
+++ b/src/io.jl
@@ -20,27 +20,33 @@ function create_parameters_and_sets_from_file(input_folder::AbstractString)
     flows_intervals_df  = read_csv_with_schema(fillpath("flows-time-intervals.csv"), TimeIntervalsData)
 
     # Sets and subsets that depend on input data
-    assets            = assets_data_df[assets_data_df.active.==true, :].name         #assets in the energy system that are active
-    assets_producer   = assets_data_df[assets_data_df.type.=="producer", :].name  #producer assets in the energy system
-    assets_consumer   = assets_data_df[assets_data_df.type.=="consumer", :].name  #consumer assets in the energy system
-    assets_storage    = assets_data_df[assets_data_df.type.=="storage", :].name  #storage assets in the energy system
-    assets_hub        = assets_data_df[assets_data_df.type.=="hub", :].name  #hub assets in the energy system
-    assets_conversion = assets_data_df[assets_data_df.type.=="conversion", :].name  #conversion assets in the energy system
-    assets_investment = assets_data_df[assets_data_df.investable.==true, :].name  #assets with investment method in the energy system
-    rep_periods       = unique(assets_profiles_df.rep_period_id)  #representative periods
-    time_steps        = Dict(row.id => 1:row.num_time_steps for row in eachrow(rep_period_df))   #time steps in the RP (e.g., hours), that are dependent on RP
-    flows             = [(row.from_asset, row.to_asset) for row in eachrow(flows_data_df)]
-    # time_intervals_per_asset = compute_time_intervals(assets_intervals_df, assets, time_steps)
-    time_intervals_per_flow = compute_time_intervals(flows_intervals_df, flows, time_steps)
+    assets                   = assets_data_df[assets_data_df.active.==true, :].name         #assets in the energy system that are active
+    assets_producer          = assets_data_df[assets_data_df.type.=="producer", :].name  #producer assets in the energy system
+    assets_consumer          = assets_data_df[assets_data_df.type.=="consumer", :].name  #consumer assets in the energy system
+    assets_storage           = assets_data_df[assets_data_df.type.=="storage", :].name  #storage assets in the energy system
+    assets_hub               = assets_data_df[assets_data_df.type.=="hub", :].name  #hub assets in the energy system
+    assets_conversion        = assets_data_df[assets_data_df.type.=="conversion", :].name  #conversion assets in the energy system
+    assets_investment        = assets_data_df[assets_data_df.investable.==true, :].name  #assets with investment method in the energy system
+    rep_periods              = unique(assets_profiles_df.rep_period_id)  #representative periods
+    time_steps               = Dict(row.id => 1:row.num_time_steps for row in eachrow(rep_period_df))   #time steps in the RP (e.g., hours), that are dependent on RP
+    flows                    = [(row.from_asset, row.to_asset) for row in eachrow(flows_data_df)]
+    time_intervals_per_asset = compute_time_intervals(assets_intervals_df, assets, time_steps)
+    time_intervals_per_flow  = compute_time_intervals(flows_intervals_df, flows, time_steps)
 
     # From balance equations:
     # Every asset a ∈ A and every rp ∈ RP will define a collection of flows, and therefore the time steps
     # can be defined a priori.
     constraints_time_periods = Dict(
         (a, rp) => begin
-            compute_rp_periods([
-                time_intervals_per_flow[(f, rp)] for f in flows if f[1] == a || f[2] == a
-            ])
+            compute_rp_periods(
+                [
+                    [
+                        time_intervals_per_flow[(f, rp)] for
+                        f in flows if f[1] == a || f[2] == a
+                    ]
+                    [time_intervals_per_asset[(a, rp)]]
+                ],
+            )
         end for a in assets, rp in rep_periods
     )
 
@@ -146,7 +152,7 @@ function create_parameters_and_sets_from_file(input_folder::AbstractString)
         flows = flows,
         rep_periods = rep_periods,
         time_steps = time_steps,
-        # time_intervals_per_asset = time_intervals_per_asset,
+        time_intervals_per_asset = time_intervals_per_asset,
         time_intervals_per_flow = time_intervals_per_flow,
         constraints_time_periods = constraints_time_periods,
     )
diff --git a/src/model.jl b/src/model.jl
index 2062cbe9..f98b5e67 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -19,10 +19,10 @@ function create_model(graph, params, sets; verbose = false, write_lp_file = fals
     Fi  = [f for f ∈ F if params.flows_investable[f]]
     Ft  = [f for f ∈ F if params.flows_is_transport[f]]
     # K_rp = sets.time_steps
-    # K_A  = sets.time_intervals_per_asset
+    K_A = sets.time_intervals_per_asset
     K_F = sets.time_intervals_per_flow
-    P   = sets.constraints_time_periods
-    RP  = sets.rep_periods
+    P = sets.constraints_time_periods
+    RP = sets.rep_periods
 
     # Model
     model = Model(HiGHS.Optimizer)
@@ -32,7 +32,7 @@ function create_model(graph, params, sets; verbose = false, write_lp_file = fals
     @variable(model, flow[f ∈ F, rp ∈ RP, K_F[(f, rp)]])         #flow from asset a to asset aa [MW]
     @variable(model, 0 ≤ assets_investment[Ai], Int)  #number of installed asset units [N]
     @variable(model, 0 ≤ flows_investment[Fi], Int)
-    @variable(model, 0 ≤ storage_level[a ∈ As, rp ∈ RP, P[(a, rp)]])
+    @variable(model, 0 ≤ storage_level[a ∈ As, rp ∈ RP, K_A[(a, rp)]])
 
     # TODO: Fix storage_level[As, RP, 0] = 0
 
@@ -95,12 +95,17 @@ function create_model(graph, params, sets; verbose = false, write_lp_file = fals
 
     # - storage balance equation
     # TODO: Add p^{inflow}
-    # TODO: Fix the initial storage_level
     @constraint(
         model,
-        c_storage_balance[a ∈ As, rp ∈ RP, (k, T) ∈ enumerate(P[(a, rp)])],
-        storage_level[a, rp, T] ==
-        (k ≥ 2 ? storage_level[a, rp, P[(a, rp)][k-1]] : 0.0) + sum(
+        c_storage_balance[a ∈ As, rp ∈ RP, T ∈ P[(a, rp)]],
+        sum(
+            ratio([T[end]], I, rp) * storage_level[a, rp, I] for
+            I ∈ sets.time_intervals_per_asset[(a, rp)]
+        ) ==
+        sum(
+            ratio([T[1] - 1], I, rp) * storage_level[a, rp, I] for
+            I ∈ sets.time_intervals_per_asset[(a, rp)]
+        ) + sum(
             ratio(T, I, rp) * flow[f, rp, I] * params.flows_efficiency[f] for
             f ∈ F, I ∈ sets.time_intervals_per_flow[(f, rp)] if f[2] == a
         ) - sum(
@@ -230,8 +235,10 @@ function create_model(graph, params, sets; verbose = false, write_lp_file = fals
     @constraint(
         model,
         upper_bound_for_storage_level[a ∈ As, rp ∈ RP, T ∈ P[(a, rp)]],
-        storage_level[a, rp, T] ≤
-        params.initial_storage_capacity[a] + (a ∈ Ai ? energy_limit[a] : 0.0)
+        sum(
+            ratio(T, I, rp) * storage_level[a, rp, I] for
+            I ∈ sets.time_intervals_per_asset[(a, rp)]
+        ) ≤ params.initial_storage_capacity[a] + (a ∈ Ai ? energy_limit[a] : 0.0)
     )
 
     if write_lp_file