import numpy as np
import matplotlib.pyplot as plt
from shapely.geometry import Point
from geopandas import GeoDataFrame

%matplotlib inline

# initial conditions
xmi = 575000   # leftmost corner of grid (m)
ymi= 4710000   # bottommost corner of grid (m)
SL = 600       # shot line interval (m)
RL = 600       # receiver line interval (m)
si = 100       # source point interval (m)
ri = 100       # receiver interval (m)
x = 3000       # width of survey (m)
y = 1800       # height of survey (m)

# Number of receiver lines and source lines
rlines = int(y/RL) + 1
slines = int(x/SL) + 1

# Put recevier lines East-West, and shot lines North South
rperline = int(x/ri) + 2
sperline = int(y/si) + 2

# offset the receivers relative to the sources
shiftx = -si/2.
shifty = -ri/2.

[x**2 for x in range(10)]

[x**y for x in range(4) for y in range(4)]

# Find x,y coordinates of recs and shots
rcvrx = [xmi + rcvr*ri + shiftx for line in range(rlines) for rcvr in range(rperline)]
rcvry = [ymi + line*RL - shifty for line in range(rlines) for rcvr in range(rperline)]

srcx = [xmi + line*SL for line in range(slines) for src in range(sperline)]
srcy = [ymi + src*si for line in range(slines) for src in range(sperline)] 

def plot_geoms(xcoords, ycoords, color = 'none', size = 50, alpha = 0.5):
    """
    A helper function to make it a bit easier to plot multiple things.
    """
    plot = plt.scatter(xcoords,
                       ycoords, 
                       c=color, 
                       s=size,
                       marker='o',
                       alpha=alpha,
                       edgecolor='none'
                       )
    return plot

fig = plt.figure(figsize=(15,10))
r = plot_geoms(rcvrx, rcvry, 'b')
s = plot_geoms(srcx, srcy, 'r')
plt.axis('equal')
plt.grid()

# Zip into x,y coordinate pairs.
rcvrxy = zip(rcvrx, rcvry)
srcxy = zip(srcx, srcy)

# Create a list of Point objects.
rcvrs = [Point(x,y) for x,y in rcvrxy]
srcs = [Point(x,y) for x,y in srcxy]

# Put lists into a GeoDataFrame.
receivers = GeoDataFrame({'geometry': rcvrs})
sources = GeoDataFrame({'geometry': srcs})

sources.to_file('sources.shp')
receivers.to_file('receivers.shp')

from shapely.geometry import LineString

midpoints = [LineString([r, s]).interpolate(0.5, normalized=True) for r in rcvrs for s in srcs]

from geopandas import GeoSeries

m = GeoSeries(midpoints)
m.plot()