from mtpy import MT, MTCollection

# Read single station
mt = MT('station001.edi')

# Access data
Z = mt.Z                         # Complex impedance tensor
freq = mt.frequency              # Frequency array
rho_xy = mt.apparent_resistivity[:, 0, 1]  # Apparent resistivity

# Station info
print(mt.station, mt.latitude, mt.longitude)

# Write EDI
mt.write_edi('output.edi')

from mtpy import MT

mt = MT('station001.edi')
print(f"Station: {mt.station}")
print(f"Location: ({mt.latitude}, {mt.longitude})")
print(f"Frequencies: {len(mt.frequency)} points")
print(f"Period range: {1/mt.frequency.max():.2f} - {1/mt.frequency.min():.0f} s")

from mtpy import MTCollection

mc = MTCollection()
mc.from_edis('survey_data/*.edi')
print(f"Loaded {len(mc)} stations")

for station in mc:
    print(f"  {station.station}: ({station.latitude:.4f}, {station.longitude:.4f})")

from mtpy import MT
from mtpy.imaging import PlotMTResponse

mt = MT('station001.edi')
plot = PlotMTResponse(mt)
plot.plot()  # Apparent resistivity and phase

from mtpy import MT
from mtpy.imaging import PlotPhaseTensor

mt = MT('station001.edi')

# Get phase tensor parameters
phi_min = mt.phase_tensor.phimin
phi_max = mt.phase_tensor.phimax
skew = mt.phase_tensor.skew       # 3D indicator

# Plot
pt = PlotPhaseTensor(mt)
pt.plot()

from mtpy import MT

mt = MT('station001.edi')
mt_rotated = mt.rotate(30)        # 30 degrees clockwise
mt.rotate_to_strike()             # Auto-rotate to geoelectric strike

from mtpy import MTCollection
from mtpy.imaging import PlotPseudoSection

mc = MTCollection()
mc.from_edis('profile/*.edi')

ps = PlotPseudoSection(mc)
ps.plot(plot_type='apparent_resistivity', mode='te')  # or 'tm', 'det'

from mtpy import MT
import pandas as pd

mt = MT('station001.edi')

# Export to CSV
df = pd.DataFrame({
    'frequency': mt.frequency,
    'rho_xy': mt.apparent_resistivity[:, 0, 1],
    'rho_yx': mt.apparent_resistivity[:, 1, 0],
    'phase_xy': mt.phase[:, 0, 1],
    'phase_yx': mt.phase[:, 1, 0]
})
df.to_csv('mt_data.csv', index=False)

# Export for ModEM
mt.write_modem('station001.dat')