import numpy as np
from numba import njit
from ode import RK4
import matplotlib.pyplot as plt

# parametry dela a strely https://en.wikipedia.org/wiki/BL_16-inch_Mk_I_naval_gun
g= 9.81            # tihove zrychleni (m/s^2)
r= 0.5*16.0*0.0254 # raze 16" (konverze na polomer a palcu na metry)
S= np.pi*r**2      # prurez strely (m^2)
rho= 1.2           # hustota vzduchu (kg/m^3)
Cd= 0.3            # cinitel odporu (bezrozmerny)
M= 929             # hmotnost strely (kg)
K= 0.5*rho*S*Cd    # pomocny koeficient
KM= K/M            # pomocny koeficient
V0= 788            # ustova rychlost strely (m/s), pevne dana
# alpha= 10*np.pi/180.0 # elevace dela (rozsah -3 az 40 stupnu, konverze na rad)
x0= 0.0            # pocatecni horizontalni poloha
z0= 10.0           # pocatecni vertikalni poloha
# Vx0= V0*np.cos(alpha) # pocatecni horizontalni rychlost
# Vz0= V0*np.sin(alpha) # pocatecni vertikalni rychlost

# skutecna vzdalenost cile - neznama
dist= 25000

@njit
def fgun(t,y):
# vypocet prave strany ODE
    (x,z,vx,vz)= y             # vektor y obsahuje postupne dve souradnice a dve rychlosti
    v= np.sqrt(vx**2 + vz**2)  # velikost rychlosti pro vypocet aerodynamicke sily
# dx/dt=  vx       # zavedeni rychlosti pro snizeni radu ODE
# dz/dt=  vz
# dvx/dt=   -KM*v*vx # horizontalni sila pouze aerodynamicka
# dvz/dt= -g-KM*v*vz # vertikalni sila vcetne gravitace
    return np.array((vx,vz,-KM*v*vx,-g-KM*v*vz))


@njit
def shoot(alpha):
# jeden vystrel se zadanou elevaci alpha (ve stupnich)
# vraci vzdalenost dopadu v metrech
    alphar= np.pi*alpha/180.0 # stupne na radiany
    Vx0= V0*np.cos(alphar) # pocatecni horizontalni rychlost
    Vz0= V0*np.sin(alphar) # pocatecni vertikalni rychlost

    tmin= 0
    tmax= 120    # maximalni cas v sekundach, dve minuty by mely stacit
    nmax= 10000  # pocet casovych kroku

    t= np.linspace(tmin,tmax,nmax+1) # ekvidistantne vzorkovany cas
    y0= np.array([x0,z0,Vx0,Vz0])    # pocatecni vektor

    y= RK4(fgun,y0,t) # integruj vystrel

# najdi misto dopadu (vertikalni souradnice projde nulou)
    n= 1
    while (y[n,1] >= 0):
        n+=1

# vykresleni trajektorie strely (nutno zakomentovat @njit pro tuto funkci!!!)
#    plt.plot(y[:n,0],y[:n,1])
#    plt.show()

# vrati horizontalni souradnici odpovidajici vzdalenosti dopadu v metrech
    return y[n,0]

@njit
def observer(l):
# Vraci True pokud je vystrel kratky (dopad strely je videt pred cilem)
# Vraci False pokud je vystrel dlouhy (dopad strely je videt za cilem)

   return True if l<dist else False

# priklad jednoho vystrelu pod uhlem 40 stupnu
alpha=40.0
l= shoot(alpha)
print('Dostrel pro uhel ',alpha,' je ',l, ' m')
if observer(l):
    print('Kratky')
else:
    print('Dlouhy')