import mplstereonet
import matplotlib.pyplot as plt

# Create stereonet
fig, ax = mplstereonet.subplots()

# Plot plane and pole (strike/dip, right-hand rule)
ax.plane(315, 45, 'b-')          # Great circle
ax.pole(315, 45, 'ko')           # Pole to plane

# Plot lineation (trend/plunge)
ax.line(120, 30, 'r^')

ax.grid()
plt.savefig('stereonet.png', dpi=150)

import numpy as np

strikes = [45, 52, 38, 48, 55, 41, 50, 43]
dips = [25, 30, 22, 28, 35, 24, 32, 27]

fig, ax = mplstereonet.subplots()

# Density contour of poles
ax.density_contourf(strikes, dips, measurement='poles', cmap='Reds')
ax.pole(strikes, dips, 'k.', markersize=5)

ax.grid()
ax.set_title('Bedding Orientations')
plt.savefig('density.png', dpi=150)

# Fit best-fit plane (girdle)
mean_strike, mean_dip = mplstereonet.fit_girdle(strikes, dips)

# Or calculate mean pole for clustered data
lon, lat = mplstereonet.pole(strikes, dips)
mean_lon, mean_lat = mplstereonet.find_mean_vector(lon, lat)
mean_s, mean_d = mplstereonet.pole2strike(mean_lon, mean_lat)

# Bedding measurements around a fold
strikes = np.array([20, 35, 50, 70, 90, 110, 130, 150, 165, 180])
dips = np.array([45, 40, 35, 30, 25, 30, 35, 40, 45, 50])

fig, ax = mplstereonet.subplots()
ax.pole(strikes, dips, 'ko', markersize=6)

# Fit girdle to poles - fold axis is pole to girdle
girdle_strike, girdle_dip = mplstereonet.fit_girdle(strikes, dips)
ax.plane(girdle_strike, girdle_dip, 'r-', linewidth=2)

fold_trend, fold_plunge = mplstereonet.pole(girdle_strike, girdle_dip)
ax.line(fold_trend, fold_plunge, 'r^', markersize=12, label='Fold axis')

ax.grid()
ax.legend()

fault_strike, fault_dip = 45, 60
rake = 30  # Degrees from strike

# Convert rake to trend/plunge
slip_trend, slip_plunge = mplstereonet.rake(fault_strike, fault_dip, rake)

fig, ax = mplstereonet.subplots()
ax.plane(fault_strike, fault_dip, 'r-', linewidth=2)
ax.line(slip_trend, slip_plunge, 'r>', markersize=10)
ax.grid()

set1 = {'strikes': [45, 50, 42, 48], 'dips': [70, 75, 68, 72]}
set2 = {'strikes': [135, 140, 130, 138], 'dips': [60, 65, 58, 62]}

fig, ax = mplstereonet.subplots()
ax.pole(set1['strikes'], set1['dips'], 'ro', label='Set 1')
ax.pole(set2['strikes'], set2['dips'], 'bs', label='Set 2')
ax.grid()
ax.legend()

# Strike/dip to dip direction
strike, dip = 45, 60
dip_direction = (strike + 90) % 360

# Pole to strike/dip
lon, lat = mplstereonet.pole(strike, dip)
back_strike, back_dip = mplstereonet.pole2strike(lon, lat)