# !pip install phiflow

# from phi.tf.flow import *
from phi.torch.flow import *
# from phi.jax.stax.flow import *

"""
Explanation: Learn to Throw

In this notebook, we will train a fully-connected neural network to solve an inverse ballistics problem.
We will compare supervised training to differentiable physics training, both numerically and visually.
Φ-Flow
&nbsp;&nbsp;&nbsp; Documentation
&nbsp;&nbsp;&nbsp; API
&nbsp;&nbsp;&nbsp; Demos
End of explanation
"""


def simulate_hit(pos, height, vel, angle, gravity=1.):
    vel_x, vel_y = math.cos(angle) * vel, math.sin(angle) * vel
    height = math.maximum(height, .5)
    hit_time = (vel_y + math.sqrt(vel_y**2 + 2 * gravity * height)) / gravity
    return pos + vel_x * hit_time, hit_time, height, vel_x, vel_y

simulate_hit(10, 1, 1, 0)[0]

"""
Explanation: To define the physics problem, we write a function to simulate the forward physics. This function takes the initial position, height, speed and angle of a thrown object and computes where the object will land, assuming the object follows a parabolic trajectory. Neglecting friction, can compute this by solving a quadratic equation.
End of explanation
"""


def sample_trajectory(pos, height, vel, angle, gravity=1.):
    hit, hit_time, height, vel_x, vel_y = simulate_hit(pos, height, vel, angle, gravity)
    def y(x):
        t = (x.vector[0] - pos) / vel_x
        y_ = height + vel_y * t - gravity / 2 * t ** 2
        return math.where((y_ > 0) & (t > 0), y_, NAN)
    return CenteredGrid(y, x=2000, bounds=Box(x=(min(pos.min, hit.min), max(pos.max, hit.max))))

vis.plot(sample_trajectory(tensor(10), 1, 1, math.linspace(-PI/4, 1.5, 7)), title="Varying Angle")

"""
Explanation: Let's plot the trajectory! We define y(x) and sample it on a grid. We have to drop invalid values, such as negative flight times and below-ground points.
End of explanation
"""


vel = 1

def loss_function(vel):
  return math.l2_loss(simulate_hit(10, 1, vel, 0)[0] - 0)

loss_function(0)

"""
Explanation: Before we train neural networks on this problem, let's perform a classical optimization using gradient descent in the initial velocity vel. We need to define a loss function to optimize. Here we desire the object to hit at x=0.
End of explanation
"""


gradient = math.functional_gradient(loss_function)
gradient(0)

"""
Explanation: Φ<sub>Flow</sub> uses the selected library (PyTorch/TensorFlow/Jax) to derive analytic derivatives.
By default, the gradient function also returns the function value.
End of explanation
"""


trj = [vel]
for i in range(10):
  loss, (grad,) = gradient(vel)
  vel = vel - .2 * grad
  trj.append(vel)
  print(f"vel={vel} - loss={loss}")
trj = math.stack(trj, channel('opt'))
vis.plot(sample_trajectory(tensor(10), 1, trj, 0))

"""
Explanation: Now we can just subtract the gradient times a learning rate $\eta = 0.2$ until we converge.
End of explanation
"""


def generate_data(shape):
  pos = math.random_normal(shape)
  height = math.random_uniform(shape) + .5
  vel = math.random_uniform(shape)
  angle = math.random_uniform(shape) * PI/2
  return math.stack(dict(pos=pos, height=height, vel=vel, angle=angle), channel('vector'))

x_train = generate_data(batch(examples=1000))
x_test = generate_data(batch(examples=1000))
y_train = simulate_hit(*x_train.vector)[0]
y_test = simulate_hit(*x_test.vector)[0]

"""
Explanation: Next, we generate a training set and test by sampling random values.
End of explanation
"""


math.seed(0)
net_sup = dense_net(1, 4, [32, 64, 32])
net_sup

# net_sup.trainable_weights[0]  # TensorFlow
net_sup.state_dict()['linear0.weight'].flatten()  # PyTorch
# net_sup.parameters[0][0]  # Stax

"""
Explanation: Now, let's create a fully-connected neural network with three hidden layers. We can reset the seed to make the weight initialization predictable.
End of explanation
"""


math.seed(0)
net_dp = dense_net(1, 4, [32, 64, 32])
# net_dp.trainable_weights[0]  # TensorFlow
net_dp.state_dict()['linear0.weight'].flatten()  # PyTorch
# net_dp.parameters[0][0]  # Stax

"""
Explanation: For the differentiable physics network we do the same thing again.
End of explanation
"""


opt_sup = adam(net_sup)

def supervised_loss(x, y, net=net_sup):
  prediction = math.native_call(net, y)
  return math.l2_loss(prediction - x)

print(f"Supervised loss (training set): {supervised_loss(x_train, y_train)}")
print(f"Supervised loss (test set): {supervised_loss(x_test, y_test)}")

for i in range(100):
  update_weights(net_sup, opt_sup, supervised_loss, x_train, y_train)

print(f"Supervised loss (training set): {supervised_loss(x_train, y_train)}")
print(f"Supervised loss (test set): {supervised_loss(x_test, y_test)}")

"""
Explanation: Indeed, the We see that the networks were initialized identically! Alternatively, we could have saved and loaded the state.
Supervised Training
Now we can train the network. We feed the desired hit position into the network and predict a possible initial state.
For supervised training, we compare the network prediction to the ground truth from our training set.
End of explanation
"""


def physics_loss(y, net=net_dp):
  prediction = math.native_call(net, y)
  y_sim = simulate_hit(*prediction.vector)[0]
  return math.l2_loss(y_sim - y)

opt_dp = adam(net_dp)

for i in range(100):
  update_weights(net_dp, opt_dp, physics_loss, y_train)

print(f"Supervised loss (training set): {supervised_loss(x_train, y_train, net=net_dp)}")
print(f"Supervised loss (test set): {supervised_loss(x_test, y_test, net=net_dp)}")

"""
Explanation: What? That's almost no progress! Feel free to run more iterations but there is a deeper problem at work here. Before we get into that, let's train a the network again but with a differentiable physics loss function.
Training with Differentiable Physics
For the differentiable physics loss, we simulate the trajectory given the initial conditions predicted by the network. Then we can measure how close to the desired location the network got.
End of explanation
"""


print(f"Supervised network (test set): {physics_loss(y_test, net=net_sup)}")
print(f"Diff.Phys. network (test set): {physics_loss(y_test, net=net_dp)}")

"""
Explanation: This looks even worse! The differentiable physics network seems to stray even further from the ground truth.
Well, we're not trying to match the ground truth, though. Let's instead measure how close to the desired location the network threw the object.
End of explanation
"""


predictions = math.stack({
    "Ground Truth": x_test.examples[:4],
    "Supervised": math.native_call(net_sup, y_test.examples[:4]),
    "Diff.Phys": math.native_call(net_dp, y_test.examples[:4]),
}, channel('curves'))
vis.plot(sample_trajectory(*predictions.vector), size=(16, 4))

"""
Explanation: Now this is much more promissing. The diff.phys. network seems to hit the desired location very accurately considering it was only trained for 100 iterations. With more training steps, this loss will go down even further, unlike the supervised network.
So what is going on here? Why does the supervised network perform so poorly?
The answer lies in the problem itself. The task is multi-modal, i.e. there are many initial states that will hit the same target.
The network only gets the target position and must decide on a single initial state. With supervised training, there is no way to know which ground truth solution occurs in the test set. The best the network can do is average nearby solutions from the training set. But since the problem is non-linear, this will give only a rough guess.
The diff.phys. network completely ignores the ground truth solutions which are not even passed to the physics_loss function. Instead, it learns to hit the desired spot, which is exactly what we want.
We can visualize the difference by looking at a couple of trajectories.
End of explanation
"""


