vplot.v

module vplot

import os
import rand
import time

// Style of a plotted data series.
pub enum Style {
	lines
	points
	lines_points
	dots
	impulses
	boxes
	steps
}

fn (s Style) str() string {
	return match s {
		.lines { 'lines' }
		.points { 'points' }
		.lines_points { 'linespoints' }
		.dots { 'dots' }
		.impulses { 'impulses' }
		.boxes { 'boxes' }
		.steps { 'steps' }
	}
}

// SeriesCfg configures the appearance of one data series.
@[params]
pub struct SeriesCfg {
pub:
	label      string // legend label; empty = notitle
	style      Style = .lines_points
	color      string // named color ('red') or hex ('#1f77b4'); empty = auto
	line_width f64 = 1.5
	point_size f64 = 1.0
}

fn (c SeriesCfg) spec() string {
	mut parts := ['with ' + c.style.str()]
	if c.line_width > 0 {
		parts << 'linewidth ${c.line_width}'
	}
	if c.style in [.points, .lines_points, .dots] && c.point_size > 0 {
		parts << 'pointsize ${c.point_size}'
	}
	if c.color != '' {
		parts << 'linecolor rgb ' + quote(c.color)
	}
	if c.label != '' {
		parts << 'title ' + quote(c.label)
	} else {
		parts << 'notitle'
	}
	return parts.join(' ')
}

// FnCfg configures the appearance of a function expression series.
@[params]
pub struct FnCfg {
pub:
	label      string
	color      string
	line_width f64 = 1.5
	samples    int // if > 0, sets `set samples N` for this plot
}

// Series is one (x, y) data series.
pub struct Series {
pub:
	x   []f64
	y   []f64
	cfg SeriesCfg
}

// Config controls how the gnuplot process is started.
@[params]
pub struct Config {
pub:
	// interactive keeps gnuplot's default GUI terminal (qt/x11/aqua), so
	// plot commands open live windows. By default vplot is headless: it
	// selects gnuplot's no-op 'unknown' terminal, and save() still works.
	interactive bool
}

// Plot wraps a persistent gnuplot process driven through its stdin pipe.
pub struct Plot {
pub mut:
	debug bool // echo every command sent to gnuplot on stderr
mut:
	proc      os.Process
	tmp_files []string
}

// new spawns a persistent gnuplot process, headless by default.
// The binary is taken from VPLOT_GNUPLOT if set, else found on PATH.
// Pass new(interactive: true) to get live GUI plot windows.
pub fn new(cfg Config) !&Plot {
	exe := gnuplot_exe()!
	mut proc := os.new_process(exe)
	proc.set_args(['-persist'])
	proc.use_stdio_ctl = true
	proc.run()
	if !proc.is_alive() {
		return error('failed to start gnuplot executable: ${exe}')
	}
	mut p := &Plot{
		proc: proc
	}
	if !cfg.interactive {
		p.cmd('set terminal unknown')!
	}
	return p
}

fn gnuplot_exe() !string {
	if e := os.getenv_opt('VPLOT_GNUPLOT') {
		if os.is_executable(e) {
			return e
		}
		return error('VPLOT_GNUPLOT is set but not executable: ${e}')
	}
	return os.find_abs_path_of_executable('gnuplot') or {
		return error('gnuplot executable not found on PATH (or set VPLOT_GNUPLOT)')
	}
}

fn quote(s string) string {
	return "'" + s.replace("'", "''") + "'"
}

// cmd sends a raw gnuplot command (a trailing newline is added if missing).
pub fn (mut p Plot) cmd(cmd string) ! {
	line := cmd.trim_right('\n') + '\n'
	if p.debug {
		eprint('[vplot] ${line}')
	}
	p.proc.stdin_write(line)
}

// set_terminal selects a gnuplot terminal, e.g. 'dumb', 'qt', 'pngcairo'.
// vplot starts headless ('unknown' terminal) unless created with
// new(interactive: true); use this to switch terminals at any time.
pub fn (mut p Plot) set_terminal(term string) ! {
	p.cmd('set terminal ${term}')!
}

pub fn (mut p Plot) set_title(title string) ! {
	p.cmd('set title ' + quote(title))!
}

pub fn (mut p Plot) set_xlabel(label string) ! {
	p.cmd('set xlabel ' + quote(label))!
}

pub fn (mut p Plot) set_ylabel(label string) ! {
	p.cmd('set ylabel ' + quote(label))!
}

pub fn (mut p Plot) set_xrange(min f64, max f64) ! {
	p.cmd('set xrange [${min}:${max}]')!
}

pub fn (mut p Plot) set_yrange(min f64, max f64) ! {
	p.cmd('set yrange [${min}:${max}]')!
}

pub fn (mut p Plot) set_grid(on bool) ! {
	if on {
		p.cmd('set grid')!
	} else {
		p.cmd('unset grid')!
	}
}

pub fn (mut p Plot) set_log_y(on bool) ! {
	if on {
		p.cmd('set logscale y')!
	} else {
		p.cmd('unset logscale y')!
	}
}

// plot_fn plots one function expression, e.g. plot_fn('sin(x)', label: 'sine').
pub fn (mut p Plot) plot_fn(expr string, cfg FnCfg) ! {
	p.plot_fns([expr], cfg)!
}

// plot_fns plots several function expressions in one plot.
// cfg applies to all of them; per-expression colors come from gnuplot's cycle.
pub fn (mut p Plot) plot_fns(exprs []string, cfg FnCfg) ! {
	if exprs.len == 0 {
		return error('no function expressions given')
	}
	if cfg.samples > 0 {
		p.cmd('set samples ${cfg.samples}')!
	}
	mut style := 'with lines'
	if cfg.line_width > 0 {
		style += ' linewidth ${cfg.line_width}'
	}
	if cfg.color != '' {
		style += ' linecolor rgb ' + quote(cfg.color)
	}
	mut parts := []string{cap: exprs.len}
	for i, e in exprs {
		mut title := 'notitle'
		if cfg.label != '' {
			label := if exprs.len > 1 { cfg.label + ' ${i + 1}' } else { cfg.label }
			title = 'title ' + quote(label)
		}
		parts << '${e} ${style} ${title}'
	}
	p.cmd('plot ' + parts.join(', '))!
}

// plot3_fn plots a 3D surface for an expression in x and y, e.g. 'sin(x)*cos(y)'.
pub fn (mut p Plot) plot3_fn(expr string, label string) ! {
	p.cmd('set hidden3d')!
	p.cmd('set pm3d')!
	title := if label != '' { 'title ' + quote(label) } else { 'notitle' }
	p.cmd('splot ${expr} with pm3d ${title}')!
}

// plot_xy plots a single (x, y) series.
pub fn (mut p Plot) plot_xy(x []f64, y []f64, cfg SeriesCfg) ! {
	p.plot_series([Series{
		x:   x
		y:   y
		cfg: cfg
	}])!
}

// plot_series plots several (x, y) series in one plot,
// using a temporary data file with gnuplot index blocks.
pub fn (mut p Plot) plot_series(series []Series) ! {
	if series.len == 0 {
		return error('no series given')
	}
	path := p.write_data_file(series)!
	mut parts := []string{cap: series.len}
	for i, s in series {
		parts << quote(path) + ' index ${i} using 1:2 ' + s.cfg.spec()
	}
	p.cmd('plot ' + parts.join(', '))!
}

fn (mut p Plot) write_data_file(series []Series) !string {
	path := os.join_path(os.temp_dir(), 'vplot_${os.getpid()}_${rand.ulid()}.dat')
	mut out := []string{}
	for i, s in series {
		if s.x.len == 0 || s.x.len != s.y.len {
			return error('series ${i}: x and y must be non-empty and of equal length')
		}
		for j in 0 .. s.x.len {
			out << '${s.x[j]} ${s.y[j]}'
		}
		if i < series.len - 1 {
			// two blank lines separate gnuplot index blocks
			out << ''
			out << ''
		}
	}
	os.write_file(path, out.join_lines() + '\n')!
	p.tmp_files << path
	return path
}

// save re-renders the current plot into path; the terminal is inferred
// from the extension: .png, .svg, .pdf, .eps, .txt (ASCII art).
// width/height only apply to raster (.png) and ASCII (.txt) output.
pub fn (mut p Plot) save(path string, width int, height int) ! {
	ext := os.file_ext(path).to_lower()
	term := match ext {
		'.png' { 'pngcairo size ${width},${height} fontscale 1' }
		'.svg' { 'svg' }
		'.pdf' { 'pdfcairo' }
		'.eps' { 'postscript eps color' }
		'.txt' { 'dumb size ${width / 10},${height / 20}' }
		else { return error('unsupported output extension "${ext}" (use .png/.svg/.pdf/.eps/.txt)') }
	}
	p.cmd('set terminal ${term}')!
	p.cmd('set output ' + quote(path))!
	p.cmd('replot')!
	p.cmd('set output')!
}

// close quits gnuplot, removes temporary data files, and waits for the process.
pub fn (mut p Plot) close() {
	p.cmd('quit') or {}
	// Give gnuplot a moment to process 'quit' before touching the pipes;
	// closing them first can leave GUI-terminal gnuplot builds hanging.
	for _ in 0 .. 100 {
		if !p.proc.is_alive() {
			break
		}
		time.sleep(20 * time.millisecond)
	}
	if p.proc.is_alive() {
		p.proc.signal_kill()
	}
	p.proc.close()
	p.proc.wait()
	for f in p.tmp_files {
		os.rm(f) or {}
	}
}